Backend pytorch: Update the code to use GPU by default (#1530)

lululxvi · Oct 30, 2023 · f0d3abe · f0d3abe
1 parent 4b8ee44
commit f0d3abe
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/deepxde/backend/pytorch/tensor.py b/deepxde/backend/pytorch/tensor.py
@@ -16,7 +16,10 @@
 # An alternative way is to use GPU by default if GPU is available, which is similar to
 # TensorFlow.
 if torch.cuda.is_available():
-    torch.set_default_tensor_type(torch.cuda.FloatTensor)
+    if Version(torch.__version__) >= Version("2.1.0"):
+        torch.set_default_device("cuda")
+    else:
+        torch.set_default_tensor_type(torch.cuda.FloatTensor)
 
 
 lib = torch