This is a simple model based on model parallelism which runs on gpu’s 0 and 1. How to save the model after training and load it back so that I can test my model on cpu.
class Net(nn.Module):
    def __init__(self, gpu0, gpu1):
        super(Net, self).__init__()
        self.gpu0 = gpu0
        self.gpu1 = gpu1
        self.conv = nn.Sequential( nn.Conv2d(1, 32, 3, 1), nn.ReLU(), nn.Conv2d(32, 64, 3, 1), nn.ReLU(),
                                 nn.MaxPool2d(2), nn.Dropout(0.25), nn.Flatten(1),
                                 ).cuda(gpu0)
        self.feat = nn.Sequential(nn.Linear(9216, 128), nn.BatchNorm1d(128),
                                  nn.ReLU(), nn.Dropout2d(0.5), nn.Linear(128, 10)
                                 ).cuda(gpu1)
        
        

    def forward(self, x):
        x = self.conv(x).cuda(self.gpu1)
        x = self.feat(x)
        output = F.log_softmax(x, dim=1)
        return output