fix: fix 0 stride in CG pool management (#291)

ELS-RD · Feb 17, 2023 · 4e4cf69 · 4e4cf69
1 parent 307e827
commit 4e4cf69
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/kernl/optimizer/pool_cuda_graphs.py b/src/kernl/optimizer/pool_cuda_graphs.py
@@ -27,9 +27,11 @@ def copy_to_pool(self, t: torch.Tensor) -> torch.Tensor:
         # 64 bits alignment
         tensor_aligned_size = get_aligned_size(t)
         new_offset = self.offset + tensor_aligned_size
+        # removes 0s from stride
+        stride_fixed = tuple(i if i > 0 else 1 for i in t.stride())
         # offset is expressed in t.dtype number of elements
         new_t = torch.as_strided(
-            self.pool.view(t.dtype), size=t.size(), stride=t.stride(), storage_offset=self.offset // t.element_size()
+            self.pool.view(t.dtype), size=t.size(), stride=stride_fixed, storage_offset=self.offset // t.element_size()
         )
         new_t.copy_(t)
         self.offset = new_offset