Sorry, I might have just confused you: I meant to say that your solution for the  state_dict error should’ve been obvious and I should’ve just updated the testing code (another script) with that solution without asking.
By testing I mean it as using a trained model to make predictions on a dataset.
To recap what happens:

During the training I get no error
If interrupt the training during an epoch and than I try to resume it after a completed epoch I get the error
If I test a model whose training was either interrupted or not, the testing works

EDIT:
By using the pdb python debugger I followed the train function during resuming and found out something.
By resuming from epoch 2, after epoch 3 step 1 the gpu state is
Fri May  1 10:56:13 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.64.00    Driver Version: 440.64.00    CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 750 Ti  On   | 00000000:01:00.0  On |                  N/A |
| 34%   33C    P0     1W /  38W |   1897MiB /  2000MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      1651      G   /usr/lib/xorg/Xorg                            15MiB |
|    0      1685      G   /usr/bin/gnome-shell                          48MiB |
|    0      2214      G   /usr/lib/xorg/Xorg                            88MiB |
|    0      2347      G   /usr/bin/gnome-shell                          94MiB |
|    0     23531      C   python                                      1633MiB |
+-----------------------------------------------------------------------------+

At epoch 3 step 2 after executing atts, dets = model(images) the gpu state is
Fri May  1 10:56:33 2020       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 440.64.00    Driver Version: 440.64.00    CUDA Version: 10.2     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  GeForce GTX 750 Ti  On   | 00000000:01:00.0  On |                  N/A |
| 34%   34C    P0     3W /  38W |   1996MiB /  2000MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      1651      G   /usr/lib/xorg/Xorg                            15MiB |
|    0      1685      G   /usr/bin/gnome-shell                          48MiB |
|    0      2214      G   /usr/lib/xorg/Xorg                            88MiB |
|    0      2347      G   /usr/bin/gnome-shell                          94MiB |
|    0     23531      C   python                                      1732MiB |
+-----------------------------------------------------------------------------+

and I get the error (the same happens if I resume from epoch 1): at least 1 step of the new epoch is completed but in the next one I run out of memory.