Ref: Fix Dim.with_size()

tum-pbs · Jan 5, 2025 · 921df39 · 921df39
1 parent ac78934
commit 921df39
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/phiml/math/_shape.py b/phiml/math/_shape.py
@@ -1108,7 +1108,10 @@ def with_size(self, size, keep_item_names=True):
             return Dim(self.name, None, self.dim_type, None)
         if keep_item_names and _size_equal(self.size, size):
             return self
-        return Dim(self.name, int(size), self.dim_type, None)
+        from ._tensors import Tensor
+        if not isinstance(size, Tensor):
+            size = int(size)
+        return Dim(self.name, size, self.dim_type, None)
 
     def with_dim_size(self, dim: Union[str, 'Shape'], size: Union[int, 'math.Tensor', str, tuple, list], keep_item_names=True):
         name = dim.name if isinstance(dim, SHAPE_TYPES) else dim