Actually I have added the synchronizations, but I am not sure whether I have done it right:
t_nonzero = []
class net(nn.Module):
        ......
   def forward():
        ......
   def inferece():
        ......
        t1 = 1000*time.time()
        idx = torch.nonzero(mask.squeeze(1))
        t2 = 1000*time.time()
        t_nonzero.append(t2-t1)
        ......

if __name__ == "__main__":
     ......
     ## warmup
     for _ in range(100):
         out = net.inference(input, ratio, patchsize)
        
     torch.cuda.synchronize()
     t0 = 1000 * time.time()
     n = 1000
     for _ in range(n):
         out = net.inference(input, ratio, patchsize)
     torch.cuda.synchronize()
     t1 = 1000 * time.time()
     fps = 1000 * n / (t1 - t0)

     print('time_nonzero:',np.mean(t_nonzero[100:]))
     print('Time:',(t1-t0)/n)

Could you please help me with it? Thanks a lot.