Merge pull request #65 from Fung-Lab/vxfung-patch-4

Vxfung patch 4

matdeeplearn/trainers/base_trainer.py

@@ -145,18 +145,18 @@ def from_config(cls, config):
         else:
             rank = torch.device("cuda" if torch.cuda.is_available() else "cpu")
             local_world_size = 1
-        dataset = cls._load_dataset(config["dataset"], config["task"]["run_mode"]) if hasattr(config["dataset"], "src") else None
+        dataset = cls._load_dataset(config["dataset"], config["task"]["run_mode"]) if "src" in config["dataset"] else None
         model = cls._load_model(config["model"], config["dataset"]["preprocess_params"], dataset, local_world_size, rank)
         optimizer = cls._load_optimizer(config["optim"], model, local_world_size)
-        sampler = cls._load_sampler(config["optim"], dataset, local_world_size, rank) if hasattr(config["dataset"], "src") else None
+        sampler = cls._load_sampler(config["optim"], dataset, local_world_size, rank) if "src" in config["dataset"] else None
         data_loader = cls._load_dataloader(
             config["optim"],
             config["dataset"],
             dataset,
             sampler,
             config["task"]["run_mode"],
             config["model"]
-        ) if hasattr(config["dataset"], "src") else None
+        ) if "src" in config["dataset"] else None
 
         scheduler = cls._load_scheduler(config["optim"]["scheduler"], optimizer)
         loss = cls._load_loss(config["optim"]["loss"])

test/configs/cpu/test_predict.yml

@@ -46,6 +46,7 @@ model:
   otf_node_attr: False
   # compute gradients w.r.t to positions and cell, requires otf_edge=True  
   gradient: False
+  model_ensemble: 1
 
 optim:
   max_epochs: 5

test/configs/cpu/test_training.yml

@@ -46,6 +46,7 @@ model:
   otf_node_attr: False
   # compute gradients w.r.t to positions and cell, requires otf_edge=True  
   gradient: False
+  model_ensemble: 1
 
 optim:
   max_epochs: 5

test/scripts/cpu/utils.py

@@ -20,11 +20,11 @@ def trainer_property(config, train: bool):
 
 def assert_valid_predictions(trainer, load: str):
     try:
-        out = trainer.predict(loader=trainer.data_loader[load], split="predict", write_output=False)
+        out = trainer.predict(loader=trainer.data_loader[0][load], split="predict", write_output=False)
         assert isinstance(out["predict"][0][0], (floating, float, integer, int))
         assert isinstance(out["ids"][0][0], str)    
         if load != "predict_loader":
             assert isinstance(out["target"][0][0], (floating, float, integer, int))
     except:
         assert False   
-
+