gerazov:

I was guessing it’s because we don’t use nn or functional to define the layers, but I don’t know. If yes what is a good way to mix equations and available layers?


You can use a mix of nn or functional to define the layers. Your code doesn’t crash on my machine.
It sounds like pytorch isn’t able to push memory to your GPU. Could you let me know the following:

which pytorch version you’re using (torch.__version__)
how you installed pytorch?
what CUDA version you’re using (in a command line, do nvcc --version)
the output of your nvidia-smi