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 torch import multiprocessing as mp

class Conv2d(nn.Conv2d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.cnt = mp.Value("d", 0)
    
    def forward(self, x):
        self.cnt.value += 1
        return super().forward(x)

net = Conv2d(3, 3, 3)
net = nn.DataParallel(net).cuda()
for _ in range(10):
    x = torch.rand(3, 3, 3, 3)
    x = x.cuda()
    net(x)