RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

[amoghskanda]

Getting this error when I try to train streetsurf on seg-10061.
Config used : withmask_nolidar.240219.yaml.
preload_on_gpu is false(by default). Cache is getting loaded onto cpu.
Error on Line411(_ret.update({k: v[frame_ind, h, w].to(device)) for k,v in _ret_image_raw.items()}) in dataloader/baseloader.py.

File "/app/neuralsim/dataio/data_loader/image_loader.py", line 258, in getitem
ret = self.sample(scene_id, cam_id, cam_fi)
File "/app/neuralsim/dataio/data_loader/image_loader.py", line 213, in sample
ground_truth = self.scene_loader.get_image_and_gts(scene_id, cam_id, cam_fi, hw[..., 0], hw[..., 1], device=self.device)
File "/app/neuralsim/dataio/data_loader/base_loader.py", line 614, in get_image_and_gts
gt = self.get_image_and_metas(scene_id, cam_id, frame_ind, h, w, device=device)
File "/app/neuralsim/dataio/data_loader/base_loader.py", line 411, in get_image_and_metas
_ret.update({k: v[frame_ind, h, w].to(device) for k,v in _ret_image_raw.items()})
File "/app/neuralsim/dataio/data_loader/base_loader.py", line 411, in
_ret.update({k: v[frame_ind, h, w].to(device) for k,v in _ret_image_raw.items()})
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

I only changed the loss from MonoDepthLoss to MonoSDFLoss, nothing else. Did anybody face this issue? Thanks