Bug fix in unstack_staggered_tensor()
@@ -33,10 +33,11 @@ def _res(tensor, axis):
def unstack_staggered_tensor(tensor):
tensors = math.unstack(tensor, -1)
+ result = []
for i, dim in enumerate(math.spatial_dimensions(tensor)):
slices = [slice(None, -1) if d != dim else slice(None) for d in math.spatial_dimensions(tensor)]
- tensors[i] = math.expand_dims(tensors[i][tuple([slice(None)]+slices)], -1)
- return tensors
+ result.append(math.expand_dims(tensors[i][tuple([slice(None)]+slices)], -1))
+ return result
def stack_staggered_components(tensors):
