pthread

gistfile1.txt

import torch
import torch.nn as nn
class X(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Linear(3,4)
    def forward(self, x):
        x = next(self.parameters())

import torchvision
data = torchvision.datasets.CIFAR10(root='./', train=False,download=True)
loader = torch.utils.data.DataLoader(data,
                                          batch_size=4,
                                          shuffle=True,
                                          num_workers=4, collate_fn=lambda x: x)
x = X()
x.cuda()
x = nn.DataParallel(x)
for tmp in loader:
    x(tmp)