CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemm`

[arbenede]

I am getting the error when I run train_helmholtz.py:

Traceback (most recent call last):
  File "experiment_scripts/train_helmholtz.py", line 70, in <module>
    use_lbfgs=opt.use_lbfgs)
  File "/home/arrigo/projects/siren/training.py", line 72, in train
    model_output = model(model_input)
  File "/home/arrigo/anaconda3/envs/siren/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/arrigo/projects/siren/modules.py", line 159, in forward
    output = self.net(coords, get_subdict(params, 'net'))
  File "/home/arrigo/anaconda3/envs/siren/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/arrigo/projects/siren/modules.py", line 93, in forward
    output = self.net(coords, params=get_subdict(params, 'net'))
  File "/home/arrigo/anaconda3/envs/siren/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/arrigo/projects/siren/torchmeta/modules/container.py", line 12, in forward
    input = module(input, params=get_subdict(params, name))
  File "/home/arrigo/anaconda3/envs/siren/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/arrigo/projects/siren/torchmeta/modules/container.py", line 12, in forward
    input = module(input, params=get_subdict(params, name))
  File "/home/arrigo/anaconda3/envs/siren/lib/python3.6/site-packages/torch/nn/modules/module.py", line 550, in __call__
    result = self.forward(*input, **kwargs)
  File "/home/arrigo/projects/siren/modules.py", line 23, in forward
    output = input.matmul(weight.permute(*[i for i in range(len(weight.shape) - 2)], -1, -2))
RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)`

I am using a RTX 4060 with 8GB of RAM. Any suggestions?