Back to what AI created. I thought I could do better.

2026-02-15 16:10:27 +01:00
parent e1f37d25e3
commit 9715cd273f
1 changed files with 1 additions and 1 deletions
@@ -2375,7 +2375,7 @@ class UnifiedMathVisitor(MathExprVisitor):
    def visitShapeFunc(self, ctx):
        val = (yield ctx.expr())
        if self._is_tensor(val):
-            return val.shape.to(self.device)
+            return torch.tensor(list(val.shape), dtype=torch.long, device=self.device)
        elif self._is_list(val):
            return torch.tensor([len(val)], dtype=torch.long, device=self.device)
        else: