PyTorch入门学习课程（第4课）---- 多GPU应用

    import torch.nn as nn
    
    
    class DataParallelModel(nn.Module):
    
    def __init__(self):
        super().__init__()
        self.block1 = nn.Linear(10, 20)
    
        # wrap block2 in DataParallel
        self.block2 = nn.Linear(20, 20)
        self.block2 = nn.DataParallel(self.block2)
    
        self.block3 = nn.Linear(20, 20)
    
    def forward(self, x):
        x = self.block1(x)
        x = self.block2(x)
        x = self.block3(x)
        return x