cupy.cuda.driver.CUDADriverError: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered

Hi,
I am getting the following error while using the apex library

I am using only 1 GPU. Tried device = torch.device('cuda:0') doesn't help

    h = self.frontend(x_, device)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/content/pase/pase/models/frontend.py", line 262, in forward
    y = self.W(h)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 547, in __call__
    result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/conv.py", line 200, in forward
    self.padding, self.dilation, self.groups)
  File "/usr/local/lib/python3.6/dist-packages/apex/amp/wrap.py", line 28, in wrapper
    return orig_fn(*new_args, **kwargs)
RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR
Traceback (most recent call last):
  File "cupy/cuda/driver.pyx", line 193, in cupy.cuda.driver.moduleUnload
  File "cupy/cuda/driver.pyx", line 82, in cupy.cuda.driver.check_status
cupy.cuda.driver.CUDADriverError: CUDA_ERROR_ILLEGAL_ADDRESS: an illegal memory access was encountered
Exception ignored in: 'cupy.cuda.function.Module.__dealloc__'

Start with the traceback locations in apex/amp/wrap.py and pase/models/frontend.py at line 262, then isolate the convolution call that produces the cuDNN error. Reproduce the illegal memory access with the reported single-GPU setup; done means identifying a reproducible cause and documenting or fixing the failing path.