diff --git a/py/expression_handlers.py b/py/expression_handlers.py index 2663942..d669c0a 100644 --- a/py/expression_handlers.py +++ b/py/expression_handlers.py @@ -102,11 +102,11 @@ class FlipHandler(NormHandler): return torch.flip(tensor, (dim,)) result = tensor.detach().clone() pivot = tensor.shape[dim] // 2 - out_slice = ( + out_slice = tuple( np.s_[:] if d != dim else np.s_[pivot:] for d in range(tensor.ndim) ) - in_slice = (np.s_[:] if d != dim else np.s_[:pivot] for d in range(tensor.ndim)) - result[*out_slice] = torch.flip(tensor[*in_slice], dims=(dim,)) + in_slice = tuple(np.s_[:] if d != dim else np.s_[:pivot] for d in range(tensor.ndim)) + result[out_slice] = torch.flip(tensor[in_slice], dims=(dim,)) return result