fix tensor.real for earlier versions of PyTorch

rtqichen · Apr 6, 2023 · 64fbe9e · 64fbe9e
1 parent d6ee52b
commit 64fbe9e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchdiffeq/_impl/misc.py b/torchdiffeq/_impl/misc.py
@@ -182,7 +182,9 @@ def forward(self, t, y, *, perturb=Perturb.NONE):
         # This dtype change here might be buggy.
         # The exact time value should be determined inside the solver,
         # but this can slightly change it due to numerical differences during casting.
-        t = t.real.to(y.abs().dtype)
+        if torch.is_complex(t):
+            t = t.real
+        t = t.to(y.abs().dtype)
         if perturb is Perturb.NEXT:
             # Replace with next smallest representable value.
             t = _nextafter(t, t + 1)