| 18 | return has_gpu and gpu_count >= min_gpus |
| 19 | |
| 20 | class ToyModel(nn.Module): |
| 21 | def __init__(self): |
| 22 | super(ToyModel, self).__init__() |
| 23 | self.net1 = nn.Linear(10, 10) |
| 24 | self.relu = nn.ReLU() |
| 25 | self.net2 = nn.Linear(10, 5) |
| 26 | |
| 27 | def forward(self, x): |
| 28 | return self.net2(self.relu(self.net1(x))) |
| 29 | |
| 30 | |
| 31 | def demo_basic(rank): |