diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..983f047
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,3 @@
+pyrightconfig.json
+**/__pycache__/
+
diff --git a/configs/multiblender.gin b/configs/multiblender.gin
index e69de29..a435e04 100644
--- a/configs/multiblender.gin
+++ b/configs/multiblender.gin
@@ -0,0 +1 @@
+Config.batching = 'single_image'
diff --git a/eval.py b/eval.py
index 73fedac..411dde4 100644
--- a/eval.py
+++ b/eval.py
@@ -103,7 +103,7 @@ def render_eval_fn(variables, _, rays):
 
       vis_suite = vis.visualize_suite(pred_distance, pred_acc)
 
-      if jax.host_id() != 0:  # Only record via host 0.
+      if jax.process_index() != 0:  # Only record via host 0.
         continue
       if not FLAGS.eval_once and idx == showcase_index:
         showcase_color = pred_color
diff --git a/internal/datasets.py b/internal/datasets.py
index 9489c85..923574f 100644
--- a/internal/datasets.py
+++ b/internal/datasets.py
@@ -72,7 +72,7 @@ def __init__(self, split, data_dir, config):
       raise ValueError(
           'the split argument should be either \'train\' or \'test\', set'
           'to {} here.'.format(split))
-    self.batch_size = config.batch_size // jax.host_count()
+    self.batch_size = config.batch_size // jax.process_count()
     self.batching = config.batching
     self.render_path = config.render_path
     self.start()
@@ -149,7 +149,7 @@ def _next_train(self):
       batch_rays = utils.namedtuple_map(lambda r: r[ray_indices], self.rays)
     elif self.batching == 'single_image':
       image_index = np.random.randint(0, self.n_examples, ())
-      ray_indices = np.random.randint(0, self.rays[0][0].shape[0],
+      ray_indices = np.random.randint(0, self.rays[0][image_index].shape[0],
                                       (self.batch_size,))
       batch_pixels = self.images[image_index][ray_indices]
       batch_rays = utils.namedtuple_map(lambda r: r[image_index][ray_indices],
diff --git a/internal/models.py b/internal/models.py
index 8a201a9..327a61f 100644
--- a/internal/models.py
+++ b/internal/models.py
@@ -238,7 +238,7 @@ def render_image(render_fn, rays, rng, chunk=8192):
   num_rays = height * width
   rays = utils.namedtuple_map(lambda r: r.reshape((num_rays, -1)), rays)
 
-  host_id = jax.host_id()
+  host_id = jax.process_index()
   results = []
   for i in range(0, num_rays, chunk):
     # pylint: disable=cell-var-from-loop
@@ -253,7 +253,7 @@ def render_image(render_fn, rays, rng, chunk=8192):
       padding = 0
     # After padding the number of chunk_rays is always divisible by
     # host_count.
-    rays_per_host = chunk_rays[0].shape[0] // jax.host_count()
+    rays_per_host = chunk_rays[0].shape[0] // jax.process_count()
     start, stop = host_id * rays_per_host, (host_id + 1) * rays_per_host
     chunk_rays = utils.namedtuple_map(lambda r: utils.shard(r[start:stop]),
                                       chunk_rays)
diff --git a/scripts/train_blender.sh b/scripts/train_blender.sh
index 8259979..c5484b1 100755
--- a/scripts/train_blender.sh
+++ b/scripts/train_blender.sh
@@ -20,6 +20,9 @@ EXPERIMENT=debug
 TRAIN_DIR=/Users/barron/tmp/nerf_results/$EXPERIMENT/$SCENE
 DATA_DIR=/Users/barron/data/nerf_synthetic/$SCENE
 
+# TRAIN_DIR=/home/akashsharma/tmp/nerf_results/$EXPERIMENT/$SCENE
+# DATA_DIR=/home/akashsharma/Documents/datasets/nerf/$SCENE
+
 rm $TRAIN_DIR/*
 python -m train \
   --data_dir=$DATA_DIR \
diff --git a/scripts/train_multiblender.sh b/scripts/train_multiblender.sh
index 612c11b..83804cc 100755
--- a/scripts/train_multiblender.sh
+++ b/scripts/train_multiblender.sh
@@ -18,11 +18,14 @@
 SCENE=lego
 EXPERIMENT=debug
 TRAIN_DIR=/Users/barron/tmp/nerf_results/$EXPERIMENT/$SCENE
+# TRAIN_DIR=/home/akashsharma/tmp/nerf_results/$EXPERIMENT/$SCENE
 DATA_DIR=/Users/barron/data/down4/$SCENE
+# DATA_DIR=/home/akashsharma/Documents/datasets/multiscale/$SCENE
 
 rm $TRAIN_DIR/*
 python -m train \
   --data_dir=$DATA_DIR \
   --train_dir=$TRAIN_DIR \
   --gin_file=configs/multiblender.gin \
-  --logtostderr
+  --logtostderr \
+  --gin_param="Config.batch_size = 1024"
diff --git a/train.py b/train.py
index 467e612..d4f5a3c 100644
--- a/train.py
+++ b/train.py
@@ -148,11 +148,10 @@ def main(unused_argv):
   rng = random.PRNGKey(20200823)
   # Shift the numpy random seed by host_id() to shuffle data loaded by different
   # hosts.
-  np.random.seed(20201473 + jax.host_id())
+  np.random.seed(20201473 + jax.process_index())
 
   config = utils.load_config()
-
-  if config.batch_size % jax.device_count() != 0:
+  if config.batch_size % jax.process_count() != 0:
     raise ValueError('Batch size must be divisible by the number of devices.')
 
   dataset = datasets.get_dataset('train', FLAGS.data_dir, config)
@@ -160,6 +159,7 @@ def main(unused_argv):
 
   rng, key = random.split(rng)
   model, variables = models.construct_mipnerf(key, dataset.peek())
+  print("Loaded model")
   num_params = jax.tree_util.tree_reduce(
       lambda x, y: x + jnp.prod(jnp.array(y.shape)), variables, initializer=0)
   print(f'Number of parameters being optimized: {num_params}')
@@ -208,12 +208,12 @@ def render_eval_fn(variables, _, rays):
   init_step = state.optimizer.state.step + 1
   state = flax.jax_utils.replicate(state)
 
-  if jax.host_id() == 0:
+  if jax.process_index() == 0:
     summary_writer = tensorboard.SummaryWriter(FLAGS.train_dir)
 
   # Prefetch_buffer_size = 3 x batch_size
   pdataset = flax.jax_utils.prefetch_to_device(dataset, 3)
-  rng = rng + jax.host_id()  # Make random seed separate across hosts.
+  rng = rng + jax.process_index()  # Make random seed separate across hosts.
   keys = random.split(rng, jax.local_device_count())  # For pmapping RNG keys.
   gc.disable()  # Disable automatic garbage collection for efficiency.
   stats_trace = []
@@ -224,7 +224,7 @@ def render_eval_fn(variables, _, rays):
       reset_timer = False
     lr = learning_rate_fn(step)
     state, stats, keys = train_pstep(keys, state, batch, lr)
-    if jax.host_id() == 0:
+    if jax.process_index() == 0:
       stats_trace.append(stats)
     if step % config.gc_every == 0:
       gc.collect()
@@ -232,7 +232,7 @@ def render_eval_fn(variables, _, rays):
     # Log training summaries. This is put behind a host_id check because in
     # multi-host evaluation, all hosts need to run inference even though we
     # only use host 0 to record results.
-    if jax.host_id() == 0:
+    if jax.process_index() == 0:
       if step % config.print_every == 0:
         summary_writer.scalar('num_params', num_params, step)
         summary_writer.scalar('train_loss', stats.loss[0], step)
@@ -295,7 +295,7 @@ def render_eval_fn(variables, _, rays):
       vis_suite = vis.visualize_suite(pred_distance, pred_acc)
 
       # Log eval summaries on host 0.
-      if jax.host_id() == 0:
+      if jax.process_index() == 0:
         psnr = math.mse_to_psnr(((pred_color - test_case['pixels'])**2).mean())
         ssim = ssim_fn(pred_color, test_case['pixels'])
         eval_time = time.time() - t_eval_start