Preserve gradients in pytorch as_tensor.

lululxvi · May 31, 2023 · be1a659 · be1a659
1 parent 0877e11
commit be1a659
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deepxde/backend/pytorch/tensor.py b/deepxde/backend/pytorch/tensor.py
@@ -71,6 +71,8 @@ def Variable(initial_value, dtype=None):
 
 
 def as_tensor(data, dtype=None):
+    if isinstance(data, (list, tuple)) and all(isinstance(d, torch.Tensor) for d in data):
+        data = torch.stack(data)
     if isinstance(data, torch.Tensor):
         if dtype is None or data.dtype == dtype:
             return data