Smarter linker

pymc-devs · Oct 27, 2024 · 494bcc5 · 494bcc5
1 parent 8b8e174
commit 494bcc5
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/pytensor/link/pytorch/linker.py b/pytensor/link/pytorch/linker.py
@@ -13,7 +13,11 @@ def input_filter(self, inp: Any) -> Any:
         return pytorch_typify(inp)
 
     def output_filter(self, var: Variable, out: Any) -> Any:
-        return out.cpu()
+        import torch
+
+        if torch.is_tensor(out):
+            return out.cpu()
+        return out
 
     def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs):
         from pytensor.link.pytorch.dispatch import pytorch_funcify