(self, x)
| 263 | self.fc2 = nn.Linear(128, 10) |
| 264 | |
| 265 | def forward(self, x): |
| 266 | if self.fused: |
| 267 | x = self.convbn1(x) |
| 268 | else: |
| 269 | x = self.conv1(x) |
| 270 | x = self.bn1(x) |
| 271 | F.relu_(x) |
| 272 | if self.fused: |
| 273 | x = self.convbn2(x) |
| 274 | else: |
| 275 | x = self.conv2(x) |
| 276 | x = self.bn2(x) |
| 277 | F.relu_(x) |
| 278 | x = F.max_pool2d(x, 2) |
| 279 | F.relu_(x) |
| 280 | x = x.flatten(1) |
| 281 | x = self.fc1(x) |
| 282 | x = self.dropout(x) |
| 283 | F.relu_(x) |
| 284 | x = self.fc2(x) |
| 285 | output = F.log_softmax(x, dim=1) |
| 286 | if fused: |
| 287 | memory_allocated[0].append(torch.cuda.memory_allocated()) |
| 288 | else: |
| 289 | memory_allocated[1].append(torch.cuda.memory_allocated()) |
| 290 | return output |
| 291 | |
| 292 | def train(model, device, train_loader, optimizer, epoch): |
| 293 | model.train() |
nothing calls this directly
no outgoing calls
no test coverage detected