Hey I am facing this Cuda error: An illegal memory access encountered. I am running my program on 4 Gpus using Pytorch’s Distributed data parallel:
Traceback (most recent call last):
  File "analysis_pipelines/training_models/tabnet.py", line 150, in <module>
    distributed_tabnet(train_data_path, test_data_path, save_model_path, local_rank)
  File "analysis_pipelines/training_models/tabnet.py", line 85, in distributed_tabnet
    tabnet_model.multiprocess_fit(X_train=X_train, y_train=y_train, eval_set=eval_set, eval_metric=eval_metric, weights=1, batch_size= batch_size, patience= patience, train_dataloader=train_loader, max_epochs=max_epochs)
  File "/home/rkudupudi/rkudupudi_tf/lib/python3.6/site-packages/pytorch_tabnet/abstract_model.py", line 371, in multiprocess_fit
    self._train_epoch(train_dataloader)
  File "/home/rkudupudi/rkudupudi_tf/lib/python3.6/site-packages/pytorch_tabnet/abstract_model.py", line 633, in _train_epoch
    batch_logs = self._train_batch(X, y)
  File "/home/rkudupudi/rkudupudi_tf/lib/python3.6/site-packages/pytorch_tabnet/abstract_model.py", line 675, in _train_batch
    loss.backward()
  File "/home/rkudupudi/rkudupudi_tf/lib/python3.6/site-packages/torch/_tensor.py", line 255, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/home/rkudupudi/rkudupudi_tf/lib/python3.6/site-packages/torch/autograd/__init__.py", line 149, in backward
    allow_unreachable=True, accumulate_grad=True)  # allow_unreachable flag
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.

The program runs on a single GPU without any problems. I tried doing CUDA_LAUCH_BLOCKING=1 while starting the program from the terminal, but the program hanged after a few epochs of training. After that I was blocked from accessing the GPUs, until I restart my machine. Kindly let me know how to fix this.
My current pytorch version is 1.9.0+cu102.
Kindly let me know how to resolve this.
Thanks