fix dataloader import

Author: ilyes319

scripts/eval_configs.py

@@ -46,7 +46,7 @@ def main():
 
     z_table = tools.AtomicNumberTable([int(z) for z in args.atomic_numbers.split(',')])
 
-    data_loader = torch_geometric.data.DataLoader(
+    data_loader = torch_geometric.dataloader.DataLoader(
         dataset=[data.AtomicData.from_config(config, z_table=z_table, cutoff=args.r_max) for config in configs],
         batch_size=args.batch_size,
         shuffle=False,

scripts/run_train.py

@@ -124,13 +124,13 @@ def main() -> None:
     logging.info(z_table)
     atomic_energies = np.array([atomic_energies_dict[args.dataset][z] for z in z_table.zs])
 
-    train_loader = torch_geometric.data.DataLoader(
+    train_loader = torch_geometric.dataloader.DataLoader(
         dataset=[data.AtomicData.from_config(c, z_table=z_table, cutoff=args.r_max) for c in collections.train],
         batch_size=args.batch_size,
         shuffle=True,
         drop_last=True,
     )
-    valid_loader = torch_geometric.data.DataLoader(
+    valid_loader = torch_geometric.dataloader.DataLoader(
         dataset=[data.AtomicData.from_config(c, z_table=z_table, cutoff=args.r_max) for c in collections.valid],
         batch_size=args.batch_size,
         shuffle=False,
@@ -293,7 +293,7 @@ def main() -> None:
     logging.info('Computing metrics for training, validation, and test sets')
     logger = tools.MetricsLogger(directory=args.results_dir, tag=tag + '_eval')
     for name, subset in [('train', collections.train), ('valid', collections.valid)] + collections.tests:
-        data_loader = torch_geometric.data.DataLoader(
+        data_loader = torch_geometric.dataloader.DataLoader(
             dataset=[data.AtomicData.from_config(config, z_table=z_table, cutoff=args.r_max) for config in subset],
             batch_size=args.batch_size,
             shuffle=False,

tests/test_data.py

@@ -35,7 +35,7 @@ def test_data_loader(self):
         data1 = AtomicData.from_config(config, z_table=table, cutoff=3.0)
         data2 = AtomicData.from_config(config, z_table=table, cutoff=3.0)
 
-        data_loader = torch_geometric.data.DataLoader(
+        data_loader = torch_geometric.dataloader.DataLoader(
             dataset=[data1, data2],
             batch_size=2,
             shuffle=True,

tests/test_models.py

@@ -41,7 +41,7 @@ def test_bo_model():
 
     atomic_data = data.AtomicData.from_config(config, z_table=table, cutoff=3.0)
 
-    data_loader = torch_geometric.data.DataLoader(
+    data_loader = torch_geometric.dataloader.DataLoader(
         dataset=[atomic_data, atomic_data],
         batch_size=2,
         shuffle=True,
@@ -66,7 +66,7 @@ def test_isolated_atom():
     r_cutoff = 3.0
     isolated_data = data.AtomicData.from_config(isolated_config, z_table=table, cutoff=r_cutoff)
     atomic_data = data.AtomicData.from_config(config, z_table=table, cutoff=r_cutoff)
-    data_loader = torch_geometric.data.DataLoader(
+    data_loader = torch_geometric.dataloader.DataLoader(
         dataset=[isolated_data, atomic_data],
         batch_size=2,
         shuffle=False,

tests/test_modules.py

@@ -43,7 +43,7 @@ def test_atomic_energies():
     energies_block = AtomicEnergiesBlock(atomic_energies=np.array([1.0, 3.0]))
 
     data = AtomicData.from_config(config, z_table=table, cutoff=3.0)
-    data_loader = torch_geometric.data.DataLoader(
+    data_loader = torch_geometric.dataloader.DataLoader(
         dataset=[data, data],
         batch_size=2,
         shuffle=True,