diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py
index a67d49299..9c37ce619 100644
--- a/nerfstudio/data/datamanagers/full_images_datamanager.py
+++ b/nerfstudio/data/datamanagers/full_images_datamanager.py
@@ -332,8 +332,13 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
         return {}
 
     def get_train_rays_per_batch(self):
-        # TODO: fix this to be the resolution of the last image rendered
-        return 800 * 800
+        """Returns resolution of the image returned from datamanager."""
+        if len(self.cached_train) != 0:
+            h = self.cached_train[0]["image"].shape[0]
+            w = self.cached_train[0]["image"].shape[1]
+            return h * w
+        else:
+            return 800 * 800
 
     def next_train(self, step: int) -> Tuple[Cameras, Dict]:
         """Returns the next training batch