update device

src/lightning.py

@@ -55,7 +55,7 @@ class DDPM(pl.LightningModule):
         self.val_data_prefix = val_data_prefix
         self.batch_size = batch_size
         self.lr = lr
-        self.torch_device = torch_device
+        self.torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.include_charges = include_charges
         self.test_epochs = test_epochs
         self.n_stability_samples = n_stability_samples
@@ -81,7 +81,7 @@ class DDPM(pl.LightningModule):
             in_node_nf=in_node_nf,
             n_dims=n_dims,
             context_node_nf=context_node_nf,
-            device=torch_device,
+            device=self.torch_device,
             hidden_nf=hidden_nf,
             activation=activation,
             n_layers=n_layers,

src/linker_size_lightning.py

@@ -45,7 +45,7 @@ class SizeClassifier(pl.LightningModule):
             hidden_nf=hidden_nf,
             out_node_nf=out_node_nf,
             n_layers=n_layers,
-            device=torch_device,
+            device=self.torch_device,
             normalization=normalization,
         )