Peter H.

The `shape[1]` is on cpu, while `x0` on gpu. Hence, a possible fix is, ``` max_val = torch.tensor(shape[1]-1, dtype=x0.dtype, device=x0.device) x1 = (x0 + 1).clamp(max=max_val) max_val = torch.tensor(shape[0]-1, dtype=y0.dtype, device=y0.device)...