Solved by Eta_C in post #4 
              
              
                You are right. I found that I have to use multiprocessing. 
The following code will work but, 

it is not elegant and generic enough
If I have 2 GPUs, it will output 1 to 20
I don’t know if there are any other side effects

import torch
import torch.nn as nn
import torch.nn.functional as F
from torc…