diff --git a/.github/workflows/code_checks.yml b/.github/workflows/code_checks.yml
index a3e383a33..7ea8b9b98 100644
--- a/.github/workflows/code_checks.yml
+++ b/.github/workflows/code_checks.yml
@@ -33,7 +33,7 @@ jobs:
           pylint mattport tests \
             --max-line-length=120 \
             --generated-members=numpy.*,torch.*,cv2.*,cv.* \
-            --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W \
+            --good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts \
             --disable=too-many-arguments \
             --disable=fixme
       - name: Test with pytest
diff --git a/.vscode/settings.json b/.vscode/settings.json
index 751433657..cea078e23 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -29,7 +29,7 @@
     "python.linting.pylintArgs": [
         "--generated-members=numpy.*,torch.*,cv2.*,cv.*",
         "--max-line-length=120",
-        "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W",
+        "--good-names=e,c,i,j,x,y,lr,K,fx,fy,cx,cy,H,W,ts",
         "--disable=too-many-arguments",
         "--disable=fixme"
     ],
diff --git a/configs/config.yaml b/configs/config.yaml
index f95fc0ce2..8693a0ccb 100644
--- a/configs/config.yaml
+++ b/configs/config.yaml
@@ -1,47 +1,28 @@
-# OPTIONS FOR DIFFERENT NERF VARIANTS
-# ----
-# sampler options
-sampler:
-  asdf: null
-  num_coarse_samples: 64
-  num_importance_samples: 128
-  perturb: 1.0
-  noise_std: 1.0
-# encoding options
-encodings:
-  asdf: null
-# network options
-network:
-  asdf: null
-# renderer options
-renderer:
-  asdf: null
-# loss options
-loss:
-  asdf: null
-  type: mse
-# ----
-
 # dataset options
 dataset:
-  data_directory: /path/to/data
-  dataset_type: blender # type of dataset
-  scene: lego
-  num_workers: 4
+  data_directory: data/lego
+  dataset_type: blender
+  downscale_factor: 1.0
 
-# camera options
-camera:
-  model: pinhole
+# dataloader options
+dataloader:
+  num_workers: 0
+  shuffle: true
+  num_images_to_sample_from: 8
+  num_rays_per_batch: 1024
 
 # optimizer options
 optimizer:
   lr: 1e-4
-  lr_scheduler: null
-  momentum: null
-  weight_decay: null
 
-# trainer options
-trainer:
-  batch_size: 2 # rays per batch
-  chunk_size: 1024 # chunks per batch
-  device: cpu
\ No newline at end of file
+# machine options
+machine_config:
+  num_gpus: 1
+  num_machines: 1
+  machine_rank: 0
+  dist_url: auto
+
+# graph options
+graph:
+  # NOTE(ethan): currently setting to nothing b/c we're still designing the config format...
+  _target_: mattport.nerf.graph.nerf.NeRFGraph
\ No newline at end of file
diff --git a/mattport/nerf/dataset/utils.py b/mattport/nerf/dataset/utils.py
index ccb1a3f49..f7c02ac9f 100644
--- a/mattport/nerf/dataset/utils.py
+++ b/mattport/nerf/dataset/utils.py
@@ -10,7 +10,7 @@ import imageio
 import numpy as np
 import torch
 
-from mattport.utils.io import load_from_json
+from mattport.utils.io import get_absolute_path, load_from_json
 
 
 class DotDict(dict):
@@ -93,7 +93,7 @@ def get_dataset_inputs(data_directory: str, dataset_type: str, downscale_factor:
         DatasetInputs: The inputs needed for generating rays.
     """
     if dataset_type == "blender":
-        dataset_inputs = load_blender_data(data_directory, downscale_factor=downscale_factor)
+        dataset_inputs = load_blender_data(get_absolute_path(data_directory), downscale_factor=downscale_factor)
     else:
         raise NotImplementedError()
 
diff --git a/mattport/nerf/field/nerf.py b/mattport/nerf/field/nerf.py
index d250634d8..7c5c82a0b 100644
--- a/mattport/nerf/field/nerf.py
+++ b/mattport/nerf/field/nerf.py
@@ -3,28 +3,32 @@ Multi Layer Perceptron
 """
 from typing import Optional
 
+import torch
 from torch import nn
 from torchtyping import TensorType
+
 from mattport.nerf.field_modules.encoding import NeRFEncoding
 from mattport.nerf.field_modules.field_outputs import DensityFieldOutput, RGBFieldOutput
 from mattport.nerf.field_modules.mlp import MLP
+from mattport.structures.rays import RayBundle, RaySamples
 
 
-class NeRFField():
+class NeRFField(nn.Module):
     """Multilayer perceptron"""
 
     def __init__(self) -> None:
+        super().__init__()
         self.build_nn_modules()
 
     def build_nn_modules(self) -> None:
-        """Initialize mulilayer perceptron."""
+        """Initialize modules."""
         self.encoding_xyz = NeRFEncoding(in_dim=3, num_frequencies=10, min_freq_exp=0.0, max_freq_exp=8.0)
         self.encoding_dir = NeRFEncoding(in_dim=3, num_frequencies=6, min_freq_exp=0.0, max_freq_exp=4.0)
         self.mlp_base = MLP(
             in_dim=self.encoding_xyz.get_out_dim(), out_dim=64, num_layers=8, layer_width=64, activation=nn.ReLU()
         )
         self.mlp_rgb = MLP(
-            in_dim=self.base_mlp.get_out_dim() + self.encoding_dir.get_out_dim(),
+            in_dim=self.mlp_base.get_out_dim() + self.encoding_dir.get_out_dim(),
             out_dim=64,
             num_layers=2,
             layer_width=64,
@@ -33,20 +37,19 @@ class NeRFField():
         self.field_output_rgb = RGBFieldOutput(in_dim=self.mlp_rgb.get_out_dim())
         self.field_output_density = DensityFieldOutput(in_dim=self.mlp_base.get_out_dim())
 
-    def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
-        """Process input with a multilayer perceptron.
+    def forward(self, xyz: TensorType["num_rays", 3], dirs: TensorType["num_rays", 3]):
+        """
 
         Args:
-            in_tensor (TensorType[..., "in_dim]): Network input
-
-        Returns:
-            TensorType[..., "out_dim"]: Network output
+            xyz: ()
+        # TODO(ethan): change the input to be something more abstracted
+        e.g., a FieldInput structure
         """
 
         encoded_xyz = self.encoding_xyz(xyz)
-        encoded_dir = self.encoding_dir(rays.direction)
+        encoded_dir = self.encoding_dir(dirs)
         base_mlp_out = self.mlp_base(encoded_xyz)
         rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
         field_rgb_output = self.field_output_rgb(rgb_mlp_out)
         field_density_out = self.field_output_density(base_mlp_out)
-        return rgb_coarse
+        return field_rgb_output, field_density_out
diff --git a/mattport/nerf/field_modules/base.py b/mattport/nerf/field_modules/base.py
index 521db07a7..d29c3ec99 100644
--- a/mattport/nerf/field_modules/base.py
+++ b/mattport/nerf/field_modules/base.py
@@ -19,7 +19,7 @@ class FieldModule(nn.Module):
     def build_nn_modules(self) -> None:
         """Function instantiates any torch.nn members within the module.
         If none exist, do nothing."""
-        raise ValueError("This function shouldn't be called.")
+        pass
 
     def set_in_dim(self, in_dim: int) -> None:
         """Sets input dimension of encoding
diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py
index 751370c6b..276f56d08 100644
--- a/mattport/nerf/field_modules/encoding.py
+++ b/mattport/nerf/field_modules/encoding.py
@@ -87,7 +87,7 @@ class NeRFEncoding(Encoding):
         Returns:
             TensorType[..., "output_dim"]: Output values will be between -1 and 1
         """
-        freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies)
+        freqs = 2 ** torch.linspace(self.min_freq, self.max_freq, self.num_frequencies).to(in_tensor.device)
         scaled_inputs = 2 * torch.pi * in_tensor[..., None] * freqs  # [..., "input_dim", "num_scales"]
         scaled_inputs = scaled_inputs.view(*scaled_inputs.shape[:-2], -1)  # [..., "input_dim" * "num_scales"]
         encoded_inputs = torch.cat([torch.sin(scaled_inputs), torch.cos(scaled_inputs)], axis=-1)
diff --git a/mattport/nerf/field_modules/field_outputs.py b/mattport/nerf/field_modules/field_outputs.py
index ad138773e..09220e528 100644
--- a/mattport/nerf/field_modules/field_outputs.py
+++ b/mattport/nerf/field_modules/field_outputs.py
@@ -23,13 +23,13 @@ class FieldOutput(FieldModule):
         self.in_dim = in_dim
         self.out_dim = out_dim
         self.activation = activation
-        self.net = self.build_nn_modules()
+        self.net = None
+        self.build_nn_modules()
 
     def build_nn_modules(self) -> None:
         layers = [nn.Linear(self.in_dim, self.out_dim)]
         if self.activation:
             layers.append(self.activation)
-        print(layers)
         self.net = nn.Sequential(*layers)
 
     def forward(self, in_tensor: TensorType[..., "in_dim"]) -> TensorType[..., "out_dim"]:
diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py
index 283027e46..d22bc3dd8 100644
--- a/mattport/nerf/field_modules/mlp.py
+++ b/mattport/nerf/field_modules/mlp.py
@@ -32,12 +32,14 @@ class MLP(FieldModule):
         """
         super().__init__()
         self.in_dim = in_dim
+        assert self.in_dim > 0
         self.out_dim = out_dim
         self.num_layers = num_layers
         self.layer_width = layer_width
         self.activation = activation
         self.out_activation = out_activation
-        self.net = self.build_nn_modules()
+        self.net = None
+        self.build_nn_modules()
 
     def build_nn_modules(self) -> None:
         """Initialize mulilayer perceptron."""
diff --git a/mattport/nerf/field_modules/ray_generator.py b/mattport/nerf/field_modules/ray_generator.py
index 5e5449060..8bfedfbf7 100644
--- a/mattport/nerf/field_modules/ray_generator.py
+++ b/mattport/nerf/field_modules/ray_generator.py
@@ -45,7 +45,7 @@ class RayGenerator(nn.Module):
         camera_index = 0
         camera_class = get_camera_model(self.num_intrinsics_params)
         camera = camera_class(*self.intrinsics[camera_index].tolist())
-        image_coords = camera.get_image_coords()
+        image_coords = camera.get_image_coords().to(ray_indices.device)
         coords = image_coords[i, j]
 
         ray_bundle = camera_class.generate_rays(intrinsics=intrinsics, camera_to_world=camera_to_world, coords=coords)
diff --git a/mattport/nerf/graph/base.py b/mattport/nerf/graph/base.py
index 332534e7c..989692b22 100644
--- a/mattport/nerf/graph/base.py
+++ b/mattport/nerf/graph/base.py
@@ -1,6 +1,7 @@
 """
 The Graph module contains all trainable parameters.
 """
+from abc import abstractmethod
 import importlib
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Set
@@ -9,6 +10,8 @@ from omegaconf import DictConfig
 
 from torch import nn
 
+from mattport.nerf.dataset.utils import DatasetInputs
+
 
 @dataclass
 class GraphInputs:
@@ -34,27 +37,38 @@ class Node:
 class Graph(nn.ModuleDict):
     """_summary_"""
 
-    def __init__(self, modules_config: DictConfig) -> None:
+    def __init__(self, intrinsics=None, camera_to_world=None) -> None:
         super().__init__()
-        self.modules_config = modules_config
-        # create graph and get ordering
-        self.roots = self.construct_graph()
-        self.module_order = self.get_module_order()
+        self.intrinsics = intrinsics
+        self.camera_to_world = camera_to_world
 
-        # 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[module_name] = module(**module_dict.meta_data)
+        self.modules = None  # NOTE(ethan): I'm turning off these features for now
+        if self.modules is not None:
+            # create graph and get ordering
+            self.roots = self.construct_graph()
+            self.module_order = self.get_module_order()
 
-        # calculate input dimensions based on module dependencies
-        for root in self.roots:
-            self.get_in_dim(root)
+            # initialize graph with known input dimensions; set default in_dim to 0
+            for module_name, module_dict in self.modules.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[module_name] = module(**module_dict.meta_data)
 
-        # instantiate torch.nn members of network
-        for _, module in self.items():
-            module.build_nn_modules()
+            # 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.items():
+                module.build_nn_modules()
+
+        self.populate_modules()  # populate the modules
+
+    @abstractmethod
+    def populate_modules(self):
+        """Initializes the modules that are part of the network."""
+        pass
 
     def get_in_dim(self, curr_node: Node) -> None:
         """Dynamically calculates and sets the input dimensions of the modules based on dependency graph
@@ -68,7 +82,7 @@ class Graph(nn.ModuleDict):
             for parent_name in curr_node.parents.keys():
                 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
+            self.modules[curr_node.name].meta_data.in_dim = in_dim
 
         for child_node in curr_node.children.values():
             if not child_node.visited_in_dim:
@@ -78,14 +92,14 @@ class Graph(nn.ModuleDict):
         """Constructs a dependency graph given the module configuration
 
         Args:
-            modules_config (dict): module definitions that make up the network
+            config (dict): module definitions that make up the network
 
         Returns:
             set: all root nodes of the constructed dependency graph
         """
         processed_modules = {}
         roots = set()
-        for module_name, module_dict in self.modules_config.items():
+        for module_name, module_dict in self.modules.items():
             if not module_name in processed_modules:
                 curr_module = Node(name=module_name, children={}, parents={})
                 processed_modules[module_name] = curr_module
@@ -125,7 +139,7 @@ class Graph(nn.ModuleDict):
         """Generates a graph and determines order of module operations using topological sorting
 
         Args:
-            modules_config (dict): module definitions that make up the network
+            config (dict): module definitions that make up the network
 
         Returns:
             list: ordering of the module names that should be executed
@@ -137,20 +151,12 @@ class Graph(nn.ModuleDict):
                 self.topological_sort(root, ordering_stack)
         return ordering_stack[::-1]
 
-    def forward(self, x):
-        """_summary_
+    @abstractmethod
+    def forward(self):
+        """Forward function that needs to be overridden."""
+        pass
 
-        Args:
-            x (_type_): _description_
-        """
-        # for name in self.module_order:
-        #     x = self.modules[name](x)
-
-        ## can index into previous results using dictionary
-        # raise NotImplementedError
-        for inputs, node_name, out_name in priority_queue:
-            self.outputs[out_name] = self.modules[node_name](inputs)
-
-            # possibly delete things from output if no longer needed
-            for old_name in :
-                del self.outputs[old_name]
+    @abstractmethod
+    def get_losses(self):
+        """Computes and returns the losses."""
+        pass
\ No newline at end of file
diff --git a/mattport/nerf/graph/nerf.py b/mattport/nerf/graph/nerf.py
index c8eac50e6..bc3b5ae73 100644
--- a/mattport/nerf/graph/nerf.py
+++ b/mattport/nerf/graph/nerf.py
@@ -3,42 +3,40 @@ Implementation of vanilla nerf.
 """
 
 
+import logging
+
 import torch
 from omegaconf import DictConfig
-from torch import TensorType, nn
+from torch import nn
+from torchtyping import TensorType
 
+from mattport.nerf.dataset.utils import DatasetInputs, get_dataset_inputs
+from mattport.nerf.field.nerf import NeRFField
 from mattport.nerf.field_modules.encoding import NeRFEncoding
 from mattport.nerf.field_modules.field_outputs import DensityFieldOutput, RGBFieldOutput
 from mattport.nerf.field_modules.mlp import MLP
-from mattport.nerf.field.nerf import NeRFField
 from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
 from mattport.nerf.renderers import RGBRenderer
 from mattport.nerf.sampler import UniformSampler
-from mattport.structures.cameras import Rays
-from mattport.nerf.dataset.utils import DatasetInputs
 
 
 class NeRFGraph(Graph):
-    def __init__(self, modules_config: DictConfig, dataset_inputs: DatasetInputs) -> None:
-        super().__init__(modules_config)
-
-        self.populate_modules(dataset_inputs=dataset_inputs)
-        # TODO(ethan): build all modules here instead of in the code below
-
-    def populate_modules(self, dataset_inputs: DatasetInputs):
+    def __init__(self, intrinsics=None, camera_to_world=None) -> None:
+        super().__init__(intrinsics=intrinsics, camera_to_world=camera_to_world)
 
+    def populate_modules(self):
         # ray generator
-        self.ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)
+        self.ray_generator = RayGenerator(self.intrinsics, self.camera_to_world)
 
         # samplers
         self.sampler_uniform = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
         self.sampler_pdf = UniformSampler(near_plane=0.1, far_plane=4.0, num_samples=64)
 
         # field
-        self.coarse = NeRFField()
-        self.fine = NeRFField()
+        self.field_coarse = NeRFField()
+        self.field_fine = NeRFField()
 
         # renderers
         self.renderer_rgb = RGBRenderer()
@@ -46,34 +44,38 @@ class NeRFGraph(Graph):
         # losses
         self.rgb_loss = MSELoss()
 
-    def forward(self, batch_indices: TensorType["num_rays", 3]):
+    def forward(self, ray_indices: TensorType["num_rays", 3]):
+        """Takes in the ray indices and renders out values."""
 
-        ray_bundle = self.ray_generator.forward(batch_indices)
-        rays = Rays(origin=ray_bundle.origins, direction=ray_bundle.directions)  # TODO(ethan): consolidate this
-        ray_samples = self.sampler_uniform(rays)
+        ray_bundle = self.ray_generator.forward(ray_indices)
+        ray_samples = self.sampler_uniform(ray_bundle)
 
-        xyz = ray_samples.locations
-        deltas = ray_samples.deltas
+        xyz = ray_samples.get_positions(ray_bundle) # (num_rays, num_samples, 3)
+        num_rays, num_samples = xyz.shape[:2]
+        xyz = xyz.view(-1, 3)
+        dirs = (
+            ray_bundle.directions.unsqueeze(1).repeat(1, num_samples, 1).view(num_rays * num_samples, 3)
+        )  # (num_rays, 3)
+        
+        field_output_rgb, field_output_density = self.field_coarse(xyz=xyz, dirs=dirs)
+        # move these back to the correct shape
+        field_output_rgb = field_output_rgb.view(num_rays, num_samples, 3)
+        field_output_density = field_output_density.view(num_rays, num_samples)
 
-        encoded_xyz = self.encoding_xyz(xyz)
-        encoded_dir = self.encoding_dir(rays.direction)
-        base_mlp_out = self.mlp_base(encoded_xyz)
-        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
-        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
-        field_density_out = self.field_output_density(base_mlp_out)
-        rgb_coarse = self.renderer_rgb(rgb=field_rgb_output, density=field_density_out, deltas=deltas)
+        deltas = ray_samples.get_deltas()
+        rgb_coarse = self.renderer_rgb(rgb=field_output_rgb, density=field_output_density, deltas=deltas)
+        
+        print(field_output_rgb)
+        print(field_output_density)
+        outputs = {
+            "rgb_coarse": rgb_coarse
+        }
+        return outputs
 
-        xyz_pdf = self.sampler_pdf(rays)  # TODO(ethan): implement pdf
-        encoded_xyz_pdf = self.encoding_xyz(xyz_pdf)
-        base_mlp_out = self.mlp_base(encoded_xyz_pdf)
-        rgb_mlp_out = self.mlp_rgb(torch.cat([encoded_dir, base_mlp_out], dim=-1))
-        field_rgb_output = self.field_output_rgb(rgb_mlp_out)
-        field_density_out = self.field_output_density(base_mlp_out)
-        rgb_fine = self.renderer_rgb(rgb=field_rgb_output, density=field_density_out, deltas=deltas)
-
-        return {"rgb_coarse": rgb_coarse, "rgb_fine": rgb_fine}
-
-    def get_loss(self, batch, render_outputs):
-        rgb_loss_coarse = self.rgb_loss(batch.rgb, render_outputs["rgb_coarse"])
-        rgb_loss_fine = self.rgb_loss(batch.rgb, render_outputs["rgb_fine"])
-        return {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
+    def get_losses(self, batch, graph_outputs):
+        # batch.pixels # (num_rays, 3)
+        losses = {}
+        # rgb_loss_coarse = self.rgb_loss(batch.pixels, graph_outputs["rgb_coarse"])
+        # rgb_loss_fine = self.rgb_loss(batch.pixels, graph_outputs["rgb_fine"])
+        # losses = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine}
+        return losses
diff --git a/mattport/nerf/sampler.py b/mattport/nerf/sampler.py
index 9177cfb41..8a88a870b 100644
--- a/mattport/nerf/sampler.py
+++ b/mattport/nerf/sampler.py
@@ -7,11 +7,7 @@ import torch
 from torch import nn
 from torchtyping import TensorType
 
-from mattport.structures import cameras
-
-RaySamples = NamedTuple(
-    "RaySamples", [("locations", TensorType[..., "num_samples", 3]), ("deltas", TensorType[..., "num_samples"])]
-)
+from mattport.structures.rays import RayBundle, RaySamples
 
 
 class Sampler(nn.Module):
@@ -32,7 +28,7 @@ class Sampler(nn.Module):
         self.num_samples = num_samples
 
     def generate_samples(
-        self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int
+        self, ray_bundle: RayBundle, near_plane: float, far_plane: float, num_samples: int
     ) -> RaySamples:
         """Encodes an input tensor.
 
@@ -48,17 +44,19 @@ class Sampler(nn.Module):
         """
         raise NotImplementedError
 
-    def forward(self, camera_rays: cameras.Rays, num_samples: int) -> TensorType[..., "num_samples", 3]:
+    def forward(self, ray_bundle: RayBundle) -> TensorType[..., "num_samples", 3]:
         """Call forward"""
-        return self.Sampler(camera_rays, num_samples)
+        return self.generate_samples(
+            ray_bundle, near_plane=self.near_plane, far_plane=self.far_plane, num_samples=self.num_samples
+        )
 
 
 class UniformSampler(Sampler):
     """Base Sampler. Intended to be subclassed"""
 
     def generate_samples(
-        self, camera_rays: cameras.Rays, near_plane: float, far_plane: float, num_samples: int
-    ) -> TensorType[..., "num_samples", 3]:
+        self, ray_bundle: RayBundle, near_plane: float, far_plane: float, num_samples: int
+    ) -> RaySamples:
         """Encodes an input tensor.
 
         Args:
@@ -75,9 +73,12 @@ class UniformSampler(Sampler):
         far_plane = far_plane or self.far_plane
         num_samples = num_samples or self.num_samples
 
-        bins = torch.linspace(near_plane, far_plane, num_samples + 1)
+        bins = torch.linspace(near_plane, far_plane, num_samples + 1).to(ray_bundle.origins.device)  # (num_samples,)
 
-        positions = (bins[..., 1:] + bins[..., :-1]) / 2.0
-        deltas = bins[..., 1:] - bins[..., :-1]
-
-        return RaySamples(locations=positions, deltas=deltas)
+        num_rays = ray_bundle.origins.shape[0]
+        ray_samples = RaySamples(
+            t_min=torch.ones_like(ray_bundle.origins[:, 0]) * near_plane,
+            t_max=torch.ones_like(ray_bundle.origins[:, 0]) * far_plane,
+            ts=bins.unsqueeze(0).repeat(num_rays, 1),  # (num_rays, num_samples)
+        )
+        return ray_samples
diff --git a/mattport/nerf/trainer.py b/mattport/nerf/trainer.py
index 9d85f06db..990093c1f 100644
--- a/mattport/nerf/trainer.py
+++ b/mattport/nerf/trainer.py
@@ -1,8 +1,11 @@
 """
 Code to train model.
 """
+import pprint
+
 import torch.distributed as dist
 from omegaconf import DictConfig
+from torch import optim
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.utils.data import DataLoader
 from tqdm import tqdm
@@ -10,18 +13,20 @@ from tqdm import tqdm
 from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch
 from mattport.nerf.dataset.utils import get_dataset_inputs
 from mattport.nerf.graph.base import Graph
-from mattport.nerf.optimizer import Optimizer
+from hydra.utils import instantiate
 
 
 class Trainer:
     """Training class"""
 
-    def __init__(self, local_rank: int, world_size: int, config: DictConfig):
+    def __init__(self, config: DictConfig, local_rank: int = 0, world_size: int = 1):
+        self.config = config
         self.local_rank = local_rank
         self.world_size = world_size
-        self.config = config
         self.train_dataset = None
+        self.train_dataloader = None
         self.test_dataset = None
+        self.test_dataloader = None
         self.graph = None
         self.optimizer = None
 
@@ -40,18 +45,21 @@ class Trainer:
             num_workers=self.config.dataloader.num_workers,
             shuffle=True,
         )
-        # TODO(ethan): implement test data
+        # TODO(ethan): implement the test data
 
     def setup_graph(self):
         """_summary_"""
-        self.graph = Graph(self.config.network).to(f"cuda:{self.local_rank}")
+        dataset_inputs = get_dataset_inputs(**self.config.dataset)
+        self.graph = instantiate(
+            self.config.graph, intrinsics=dataset_inputs.intrinsics, camera_to_world=dataset_inputs.camera_to_world
+        ).to(f"cuda:{self.local_rank}")
         if self.world_size > 1:
             self.graph = DDP(self.graph, device_ids=[self.local_rank])
             dist.barrier(device_ids=[self.local_rank])
 
     def setup_optimizer(self):
         """_summary_"""
-        self.optimizer = Optimizer(params=self.graph.parameters(), **self.config.optimizer)
+        self.optimizer = optim.Adam(self.graph.parameters(), lr=self.config.optimizer.lr)
 
     def load_checkpoint(self):
         """_summary_
@@ -72,21 +80,29 @@ class Trainer:
     def train(self) -> None:
         """_summary_"""
         # is_main_thread = self.local_rank % self.world_size == 0
-        raise NotImplementedError
+        # raise NotImplementedError
+        num_epochs = 10
+        for _ in range(num_epochs):
+            self.train_epoch()
 
     def train_epoch(self):
         """_summary_"""
-        for _ in range(len(self.train_dataset)):
+        num_iters = 100
+        for _ in range(num_iters):
             batch = next(iter(self.train_dataloader))
             self.train_iteration(batch)
 
     def train_iteration(self, batch):
         """_summary_"""
         # TODO(): save checkpoints, and do logging only on rank 0 device only
-        # if gpu == 0:
-        output_dict = self.graph(batch.indices)
-        
-        loss = self.graph.get_loss(batch, )
+        if self.local_rank == 0:
+            pass
 
-        loss.backward()
-        # update model
+        # move batch to correct device
+        batch.indices = batch.indices.to(f"cuda:{self.local_rank}")
+
+        ray_indices = batch.indices
+        graph_outputs = self.graph(ray_indices)
+        losses = self.graph.get_losses(batch, graph_outputs)
+
+        # TODO(ethan): update the network parameters
diff --git a/mattport/structures/rays.py b/mattport/structures/rays.py
index 6ba272c21..81b3f427e 100644
--- a/mattport/structures/rays.py
+++ b/mattport/structures/rays.py
@@ -5,6 +5,7 @@ import random
 from dataclasses import dataclass
 from typing import Optional
 
+import torch
 from torchtyping import TensorType
 
 
@@ -65,3 +66,46 @@ class CameraRayBundle:
     origins: TensorType["image_height", "image_width", 3]
     directions: TensorType["image_height", "image_width", 3]
     camera_index: int = None
+
+
+@dataclass
+class RaySamples:
+    """_summary_
+
+    Args:
+        t_min: maximum times
+        t_max: minimum times
+        ts: time steps
+    """
+
+    t_min: TensorType["num_rays"]
+    t_max: TensorType["num_rays"]
+    ts: TensorType["num_rays", "num_samples+1"]  # TODO(ethan): check with Matt about this. 2nd dimension is num_samples+1
+
+    def get_positions(self, ray_bundle: RayBundle) -> TensorType["num_rays", "num_samples", 3]:
+        """Returns positions."""
+        # TODO(ethan): check with Matt about this
+        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2 # (num_rays, num_samples)
+        return ray_bundle.origins[:, None] + t_mid[:, :, None] * ray_bundle.directions[:, None]
+
+    def get_starts(self):
+        """_summary_
+
+        Returns:
+            _type_: _description_
+        """
+        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2
+        return torch.cat([self.t_min[:, None], t_mid], dim=-1)
+
+    def get_ends(self):
+        """_summary_
+
+        Returns:
+            _type_: _description_
+        """
+        t_mid = (self.ts[:, 1:] + self.ts[:, :-1]) / 2
+        return torch.cat([t_mid, self.t_max[:, None]], dim=-1)
+
+    def get_deltas(self) -> TensorType["num_samples"]:
+        """Returns deltas."""
+        return self.get_ends() - self.get_starts()
diff --git a/scripts/run_train_nerf.py b/scripts/run_train_nerf.py
index e0e4dbbf2..a71150041 100644
--- a/scripts/run_train_nerf.py
+++ b/scripts/run_train_nerf.py
@@ -112,7 +112,7 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
     Returns:
         Any: TODO(): determine the return type
     """
-    trainer = Trainer(local_rank, world_size, config)
+    trainer = Trainer(config, local_rank, world_size)
     trainer.setup_graph()
     trainer.setup_optimizer()
     trainer.setup_dataset()