Thank you for your answer. I will do it. Another problem related to this context is the following. I now have two of these scripts, each one will call 10 models in sequence. I am running one script in one GPU and another one in another GPU. Let us say that script/GPU1 one is running model 5 and script/GPU2 is running model 3 (I started GPU1 first). Since I save the best model during training to load later in the inference/test phase (as it is usual), a problem like that happens:
size mismatch for fc1.0.weight: copying a param with shape torch.Size([2, 3072]) from checkpoint, the shape in current model is torch.Size([2, 3232]).

It is clear to me that the problem is that, after finishing training model 5 (GPU1), it loads the model from network 3 (GPU2). I am just wondering how to handle this checkpoint issue. Do I need to resort in a parallel solution in this way?
Thank you again.