Merge pull request #5 from yhgon/patch-1

Handle different device in `torch.index_select`

Author: AminRezaei0x443

memory_efficient_attention/utils.py

@@ -6,7 +6,7 @@ def dynamic_slice(x, starts, sizes):
     # start_indices[i] = clamp(start_indices[i], 0, operand.dimension_size[i] - size_indices[i])
     starts = [np.clip(starts[i], 0, x.shape[i] - sizes[i]) for i in range(len(starts))]
     for i, (start, size) in enumerate(zip(starts, sizes)):
-        x = torch.index_select(x, i, torch.tensor(range(start, start + size)))
+        x = torch.index_select(x, i, torch.tensor(range(start, start + size), device=x.device))
     return x