ValueError on vgg16_bn unless I first train for one epoch

[austinmw]

I'm able to build and display a graph right away on the torchvision alexnet model. But on the torchvision vgg16_bn model, if try to graph the model before training I get the error, ValueError: Expected more than 1 value per channel when training, got input size torch.Size([1, 1024]) Even though I passed in hl.build_graph(model, torch.zeros((1,3,224,224).cuda()). If I train the model for 1 epoch, then I'm able to build and display the graph properly.

Edit: Actually the torchvision model copied exactly works, but if I create a fastai learner with this model (which removes the head and adds a custom head with adaptive pooling) then I get this issue.

Edit 2: if I set the model to eval mode before trying to graph it works

Any idea why this is the case?

[waleedka]

HiddenLayer is designed to take either a TensorFlow or a PyTorch model. Are you passing a FastAI object instead? That might be the issue.