def forward(self, x):
mc.save(None, x)
x = F.relu(F.max_pool2d(self.conv1(x), 2))
mc.save('layer1', x)
x = F.relu(F.max_pool2d(self.conv2(x), 2))
mc.save('layer2', x)
x = torch.flatten(x, 1)
x = F.relu(self.fc1(x))
mc.save('relu', x)
x = self.fc2(x)
x = F.log_softmax(x, dim=1)
mc.save('result', x)
return x
def forward(self, x):
mc.save(None, x)
x = F.relu(F.max_pool2d(self.conv1(x), 2))
mc.save('layer1', x)
x = F.relu(F.max_pool2d(self.conv2(x), 2))
mc.save('layer2', x)
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
mc.save('relu', x)
x = self.fc2(x)
x = F.log_softmax(x, dim=1)
mc.save('result', x)
return x