diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py
index 008920d97..57c99a900 100644
--- a/mattport/nerf/field_modules/mlp.py
+++ b/mattport/nerf/field_modules/mlp.py
@@ -39,7 +39,7 @@ class MLP(FieldModule):
         self.out_activation = out_activation
         self.net = None
 
-    def build_model(self) -> None:
+    def build_nn_modules(self) -> None:
         """Initialize mulilayer perceptron."""
         layers = []
         for i in range(self.num_layers):
diff --git a/mattport/nerf/graph.py b/mattport/nerf/graph.py
index 5dd87440d..1f951e8bd 100644
--- a/mattport/nerf/graph.py
+++ b/mattport/nerf/graph.py
@@ -30,7 +30,7 @@ class Node:
         return hash(self.name)
 
 
-class Graph(nn.Module):
+class Graph(nn.ModuleDict):
     """_summary_"""
 
     def __init__(self, modules_config: DictConfig) -> None:
@@ -40,20 +40,19 @@ class Graph(nn.Module):
         self.roots = self.construct_graph()
         self.module_order = self.get_module_order()
 
-        self.modules = {}
         # initialize graph with known input dimensions; set default in_dim to 0
         for module_name, module_dict in modules_config.items():
             module = getattr(importlib.import_module("mattport.nerf.field_modules"), module_dict.class_name)
             if not module_dict.meta_data.in_dim:
                 module_dict.meta_data.in_dim = 0
-            self.modules[module_name] = module(**module_dict.meta_data)
+            self[module_name] = module(**module_dict.meta_data)
 
         # calculate input dimensions based on module dependencies
         for root in self.roots:
             self.get_in_dim(root)
 
         # instantiate torch.nn members of network
-        for _, module in self.modules.items():
+        for _, module in self.items():
             module.build_nn_modules()
 
     def get_in_dim(self, curr_node: Node) -> None:
@@ -66,8 +65,8 @@ class Graph(nn.Module):
         if len(curr_node.parents) > 0:
             in_dim = 0
             for parent_name in curr_node.parents.keys():
-                in_dim += self.modules[parent_name].get_out_dim()
-            self.modules[curr_node.name].set_in_dim(in_dim)
+                in_dim += self[parent_name].get_out_dim()
+            self[curr_node.name].set_in_dim(in_dim)
             self.modules_config[curr_node.name].meta_data.in_dim = in_dim
 
         for child_node in curr_node.children.values():
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 0b900d6c8..de8ccafc3 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -1,6 +1,7 @@
 """
 Code to train model.
 """
+from omegaconf import DictConfig
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP
 
@@ -12,7 +13,7 @@ from mattport.nerf.optimizer import Optimizer
 class Trainer:
     """Training class"""
 
-    def __init__(self, local_rank: int, world_size: int, config: dict):
+    def __init__(self, local_rank: int, world_size: int, config: DictConfig):
         self.local_rank = local_rank
         self.world_size = world_size
         self.config = config
@@ -31,11 +32,10 @@ class Trainer:
 
     def setup_graph(self):
         """_summary_"""
-        self.graph = Graph(self.config).to(self.device)
+        self.graph = Graph(self.config.network).to(self.device)
         if self.world_size > 1:
             self.graph = DDP(self.graph, device_ids=[self.local_rank])
             dist.barrier(device_ids=[self.local_rank])
-        raise NotImplementedError
 
     def setup_optimizer(self):
         """_summary_"""
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index 54d37f099..e0e4dbbf2 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -174,7 +174,7 @@ def launch(
                 process.join()
                 logger.info("Process %s finished", str(i))
     else:
-        main_func(config)
+        main_func(local_rank=0, world_size=1, config=config)
 
 
 @hydra.main(config_path="../configs", config_name="config.yaml")
@@ -186,7 +186,7 @@ def main(config: DictConfig):
         num_machines=config.machine_config.num_machines,
         machine_rank=config.machine_config.machine_rank,
         dist_url=config.machine_config.dist_url,
-        config=(config,),
+        config=config,
     )
 
 
diff --git a/tests/nerf/graph/test_graph.py b/tests/nerf/graph/test_graph.py
index e45665341..191ae9aa8 100644
--- a/tests/nerf/graph/test_graph.py
+++ b/tests/nerf/graph/test_graph.py
@@ -32,17 +32,17 @@ def instantiate_graph(config_name: str) -> Graph:
 def test_graph_init():
     """test field modules dictionary is instantiated properly"""
     test_graph, _ = instantiate_graph("test_basic_graph")
-    assert test_graph.modules is not None
-    assert type(test_graph.modules["encoder_0"]).__name__ == "ScalingAndOffset"
-    assert type(test_graph.modules["mlp_0"]).__name__ == "MLP"
-    assert type(test_graph.modules["mlp_1"]).__name__ == "MLP"
+    assert test_graph is not None
+    assert type(test_graph["encoder_0"]).__name__ == "ScalingAndOffset"
+    assert type(test_graph["mlp_0"]).__name__ == "MLP"
+    assert type(test_graph["mlp_1"]).__name__ == "MLP"
 
 
 def test_input_dimension():
     """test field module input dimension calculation"""
     test_graph, _ = instantiate_graph("test_repeat_graph")
-    assert test_graph.modules["mlp_0"].in_dim == 16
-    assert test_graph.modules["mlp_1"].in_dim == 1
+    assert test_graph["mlp_0"].in_dim == 16
+    assert test_graph["mlp_1"].in_dim == 1
 
 
 def test_construct_graph_basic():