requirements of the weights scales quadratically with the number of
features.  It is very easy to blow through your memory this way (and
remember that you will need at least twice the size of the weights,
since you also need to store the gradients.)
Consider checkpointing. You can trade-off memory for compute by
using checkpoint.
My GPU memory isn't freed properly
PyTorch uses a caching memory allocator to speed up memory
allocations. As a result, the values shown in "nvidia-smi" usually
don't reflect the true memory usage. See Memory management for more
details about GPU memory management.
If your GPU memory isn't freed even after Python quits, it is very
likely that some Python subprocesses are still alive. You may find
them via "ps -elf | grep python" and manually kill them with "kill -9
[pid]".
My out of memory exception handler can't allocate memory