runtimeError when training model from Xenium data #95

I am trying to train the model using my own Xenium data without gene embedding. I followed this https://elihei2.github.io/segger_dev/notebooks/segger_tutorial/.

Here is the error message when starting the training (i.e running trainer.fit):

/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [96,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [97,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [98,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [99,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [100,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [101,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [102,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [103,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [104,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [105,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [106,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [107,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [108,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [109,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [110,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [111,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [112,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [113,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [114,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [115,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [116,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [117,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [118,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [119,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [120,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [121,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [122,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [123,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [124,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [125,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [126,0,0] Assertion srcIndex < srcSelectDimSize failed.
/pytorch/aten/src/ATen/native/cuda/Indexing.cu:1422: indexSelectLargeIndex: block: [122,0,0], thread: [127,0,0] Assertion srcIndex < srcSelectDimSize failed.

RuntimeError Traceback (most recent call last)
File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:48, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
47 return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
---> 48 return trainer_fn(*args, **kwargs)
50 except _TunerExitException:

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:599, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
593 ckpt_path = self._checkpoint_connector._select_ckpt_path(
594 self.state.fn,
595 ckpt_path,
596 model_provided=True,
597 model_connected=self.lightning_module is not None,
598 )
--> 599 self._run(model, ckpt_path=ckpt_path)
601 assert self.state.stopped

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1012, in Trainer._run(self, model, ckpt_path)
1009 # ----------------------------
1010 # RUN THE TRAINER
1011 # ----------------------------
-> 1012 results = self._run_stage()
1014 # ----------------------------
1015 # POST-Training CLEAN UP
1016 # ----------------------------

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1054, in Trainer._run_stage(self)
1053 with isolate_rng():
-> 1054 self._run_sanity_check()
1055 with torch.autograd.set_detect_anomaly(self._detect_anomaly):

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1083, in Trainer._run_sanity_check(self)
1082 # run eval step
-> 1083 val_loop.run()
1085 call._call_callback_hooks(self, "on_sanity_check_end")

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/loops/utilities.py:179, in _no_grad_context.._decorator(self, *args, **kwargs)
178 with context_manager():
--> 179 return loop_run(self, *args, **kwargs)

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:145, in _EvaluationLoop.run(self)
144 # run step hooks
--> 145 self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
146 except StopIteration:
147 # this needs to wrap the *_step call too (not just next) for dataloader_iter support

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/loops/evaluation_loop.py:437, in _EvaluationLoop._evaluation_step(self, batch, batch_idx, dataloader_idx, dataloader_iter)
432 step_args = (
433 self._build_step_args_from_hook_kwargs(hook_kwargs, hook_name)
434 if not using_dataloader_iter
435 else (dataloader_iter,)
436 )
--> 437 output = call._call_strategy_hook(trainer, hook_name, *step_args)
439 self.batch_progress.increment_processed()

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:328, in _call_strategy_hook(trainer, hook_name, *args, **kwargs)
327 with trainer.profiler.profile(f"[Strategy]{trainer.strategy.class.name}.{hook_name}"):
--> 328 output = fn(*args, **kwargs)
330 # restore current_fx when nested context

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py:412, in Strategy.validation_step(self, *args, **kwargs)
411 return self._forward_redirection(self.model, self.lightning_module, "validation_step", *args, **kwargs)
--> 412 return self.lightning_module.validation_step(*args, **kwargs)

File /remote/home/eloiseb/code/postdoctoral/spatial_t/segger/training/train.py:190, in LitSegger.validation_step(self, batch, batch_idx)
189 # Forward pass to get the logits
--> 190 z = self.model(batch.x_dict, batch.edge_index_dict)
191 output = self.model.decode(z, edge_label_index)

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
1738 else:
-> 1739 return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1748 or _global_backward_pre_hooks or _global_backward_hooks
1749 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750 return forward_call(*args, **kwargs)
1752 result = None

File /remote/home/eloiseb/code/postdoctoral/spatial_t/segger/models/segger_model.py:111, in Segger.forward(self, x_dict, edge_index_dict)
103 """
104 Forward pass for the Segger model.
105
(...)
108 edge_index_dict (dict[str, Tensor]): Edge indices for each edge type.
109 """
--> 111 x_dict = {key: self.node_initkey for key, x in x_dict.items()}
113 x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()}

File /remote/home/eloiseb/code/postdoctoral/spatial_t/segger/models/segger_model.py:111, in (.0)
103 """
104 Forward pass for the Segger model.
105
(...)
108 edge_index_dict (dict[str, Tensor]): Edge indices for each edge type.
109 """
--> 111 x_dict = {key: self.node_initkey for key, x in x_dict.items()}
113 x_dict = {key: F.leaky_relu(x) for key, x in x_dict.items()}

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:1739, in Module._wrapped_call_impl(self, *args, **kwargs)
1738 else:
-> 1739 return self._call_impl(*args, **kwargs)

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:1750, in Module._call_impl(self, *args, **kwargs)
1747 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
1748 or _global_backward_pre_hooks or _global_backward_hooks
1749 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1750 return forward_call(*args, **kwargs)
1752 result = None

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/sparse.py:190, in Embedding.forward(self, input)
189 def forward(self, input: Tensor) -> Tensor:
--> 190 return F.embedding(
191 input,
192 self.weight,
193 self.padding_idx,
194 self.max_norm,
195 self.norm_type,
196 self.scale_grad_by_freq,
197 self.sparse,
198 )

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/functional.py:2551, in embedding(input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse)
2550 no_grad_embedding_renorm(weight, input, max_norm, norm_type)
-> 2551 return torch.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)

RuntimeError: CUDA error: device-side assert triggered
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

During handling of the above exception, another exception occurred:

RuntimeError Traceback (most recent call last)
Cell In[6], line 1
----> 1 trainer.fit(
2 model=ls,
3 datamodule=dm
4 )

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:561, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
559 self.training = True
560 self.should_stop = False
--> 561 call._call_and_handle_interrupt(
562 self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
563 )

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:69, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
67 except BaseException as exception:
68 _interrupt(trainer, exception)
---> 69 trainer._teardown()
70 # teardown might access the stage so we reset it after
71 trainer.state.stage = None

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1035, in Trainer._teardown(self)
1032 def _teardown(self) -> None:
1033 """This is the Trainer's internal teardown, unrelated to the teardown hooks in LightningModule and Callback;
1034 those are handled by :meth:_call_teardown_hook."""
-> 1035 self.strategy.teardown()
1036 loop = self._active_loop
1037 # loop should never be None here but it can because we don't know the trainer stage with ddp_spawn

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py:536, in Strategy.teardown(self)
534 if self.lightning_module is not None:
535 log.debug(f"{self.class.name}: moving model to CPU")
--> 536 self.lightning_module.cpu()
537 self.precision_plugin.teardown()
538 assert self.accelerator is not None

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/lightning_fabric/utilities/device_dtype_mixin.py:82, in _DeviceDtypeModuleMixin.cpu(self)
80 """See :meth:torch.nn.Module.cpu."""
81 _update_properties(self, device=torch.device("cpu"))
---> 82 return super().cpu()

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:1121, in Module.cpu(self)
1112 def cpu(self: T) -> T:
1113 r"""Move all model parameters and buffers to the CPU.
1114
1115 .. note::
(...)
1119 Module: self
1120 """
-> 1121 return self._apply(lambda t: t.cpu())

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:903, in Module._apply(self, fn, recurse)
901 if recurse:
902 for module in self.children():
--> 903 module._apply(fn)
905 def compute_should_use_set_data(tensor, tensor_applied):
906 if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
907 # If the new tensor has compatible tensor type as the existing tensor,
908 # the current behavior is to change the tensor in-place using .data =,
(...)
913 # global flag to let the user control whether they want the future
914 # behavior of overwriting the existing tensor or not.

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:903, in Module._apply(self, fn, recurse)
901 if recurse:
902 for module in self.children():
--> 903 module._apply(fn)
905 def compute_should_use_set_data(tensor, tensor_applied):
906 if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
907 # If the new tensor has compatible tensor type as the existing tensor,
908 # the current behavior is to change the tensor in-place using .data =,
(...)
913 # global flag to let the user control whether they want the future
914 # behavior of overwriting the existing tensor or not.

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:903, in Module._apply(self, fn, recurse)
901 if recurse:
902 for module in self.children():
--> 903 module._apply(fn)
905 def compute_should_use_set_data(tensor, tensor_applied):
906 if torch._has_compatible_shallow_copy_type(tensor, tensor_applied):
907 # If the new tensor has compatible tensor type as the existing tensor,
908 # the current behavior is to change the tensor in-place using .data =,
(...)
913 # global flag to let the user control whether they want the future
914 # behavior of overwriting the existing tensor or not.

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:930, in Module._apply(self, fn, recurse)
926 # Tensors stored in modules are graph leaves, and we don't want to
927 # track autograd history of param_applied, so we have to use
928 # with torch.no_grad():
929 with torch.no_grad():
--> 930 param_applied = fn(param)
931 p_should_use_set_data = compute_should_use_set_data(param, param_applied)
933 # subclasses may have multiple child tensors so we need to use swap_tensors

File ~/miniconda3/envs/segger-env/lib/python3.10/site-packages/torch/nn/modules/module.py:1121, in Module.cpu..(t)
1112 def cpu(self: T) -> T:
1113 r"""Move all model parameters and buffers to the CPU.
1114
1115 .. note::
(...)
1119 Module: self
1120 """
-> 1121 return self._apply(lambda t: t.cpu())

RuntimeError: CUDA error: device-side assert triggered
Compile with TORCH_USE_CUDA_DSA to enable device-side assertions.

Click to add a cell.

Thank you for your help.