ok, by manually inserting tons of logging messages all over my code, I finally found the place where the error occurs. It is the following line
    self.mvn = MultivariateNormal(loc=torch.zeros(d, device=data.device), covariance_matrix=cov)

When this line is called with d=1, cov=tensor([[0.1121]]) and data.device = “cuda:0” I receive the error message:
CUBLAS error: out of memory (3) in magma_spotrf_LL_expert_gpu at /opt/conda/conda-bld/magma-cuda102_1583546904148/work/src/spotrf_gpu.cpp:138
CUBLAS error: not initialized (1) in magma_spotrf_LL_expert_gpu at /opt/conda/conda-bld/magma-cuda102_1583546904148/work/src/spotrf_gpu.cpp:138
The training loop continues then until it comes again to this point where it produces the same error again. After a few iteration the script crashes with the error:
python: /opt/conda/conda-bld/magma-cuda102_1583546904148/work/interface_cuda/interface.cpp:810: void magma_queue_create_internal(magma_device_t, magma_queue**, const char*, const char*, int): Assertion `queue->dCarray__ != __null’ failed.
Aborted (core dumped)
I actually can’t imagine what could be so problematic about a 1-dim gaussian variable. Does anybody have a clue?