bugfix: processor device, in-place ops and more

Author: shuyijia

matdeeplearn/common/data.py

@@ -103,6 +103,7 @@ def get_dataloader(
     batch_size: int,
     num_workers: int = 0,
     sampler=None,
+    shuffle=True
 ):
     """
     Returns a single dataloader for a given dataset
@@ -124,7 +125,7 @@ def get_dataloader(
     loader = DataLoader(
         dataset,
         batch_size=batch_size,
-        shuffle=(sampler is None),
+        shuffle=shuffle,
         num_workers=num_workers,
         sampler=sampler,
     )

matdeeplearn/preprocessor/datasets.py

@@ -1,37 +1,44 @@
-import os
+import torch, os
 
-import torch
 from torch_geometric.data import InMemoryDataset
 
-
 class StructureDataset(InMemoryDataset):
     def __init__(
         self,
-        root,
-        processed_data_path,
-        transform=None,
-        pre_transform=None,
+        root, 
+        processed_data_path, 
+        transform=None, 
+        pre_transform=None, 
         pre_filter=None,
+        device=None
     ):
         self.root = root
         self.processed_data_path = processed_data_path
-        super(StructureDataset, self).__init__(
-            root, transform, pre_transform, pre_filter
-        )
-        self.data, self.slices = torch.load(self.processed_paths[0])
-
+        super(StructureDataset, self).__init__(root, transform, pre_transform, pre_filter)
+
+        if device is None:
+            try:
+                self.data, self.slices = torch.load(self.processed_paths[0])
+            except:
+                self.data, self.slices = torch.load(self.processed_paths[0], map_location=torch.device('cpu'))
+        else:
+            if device == 'cpu':
+                self.data, self.slices = torch.load(self.processed_paths[0], map_location=torch.device(device))
+            else:
+                self.data, self.slices = torch.load(self.processed_paths[0])
+    
     @property
     def raw_file_names(self):
-        """
-        The name of the files in the self.raw_dir folder
+        '''
+        The name of the files in the self.raw_dir folder 
         that must be present in order to skip downloading.
-        """
+        '''
         return []
 
     def download(self):
-        """
+        '''
         Download required data files; to be implemented
-        """
+        '''
         pass
 
     @property
@@ -40,12 +47,11 @@ def processed_dir(self):
 
     @property
     def processed_file_names(self):
-        """
-        The name of the files in the self.processed_dir
+        '''
+        The name of the files in the self.processed_dir 
         folder that must be present in order to skip processing.
-        """
+        '''
         return ["data.pt"]
 
-
 class LargeStructureDataset(InMemoryDataset):
-    pass
+    pass