Thanks @ptrblck. This reproduces the error for me:
import os
from pytorchcv.model_provider import get_model as ptcv_get_model
import torch
import types
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
os.environ["CUDA_VISIBLE_DEVICES"] = '0,1,2,3'
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
cifar_loc = '/disk/scratch/s1874193/datasets/cifar'
net = ptcv_get_model("densenet40_k12_cifar10", root = '/home/s1874193/Distillation/xdistill/pre_trained_models', pretrained=True)
def my_forward(self, x):
activations = []
for module in self.features._modules.values():
x = module(x)
activations.append(x)
x = x.view(x.size(0), -1)
x = self.output(x)
return x, activations
net.forward = types.MethodType(my_forward, net)
if torch.cuda.device_count() > 1:
net = nn.DataParallel(net, device_ids=[0,1,2,3])
net.to(device)
net.eval()
x = torch.randn(4, 3, 32, 32)
out, act = net(x)
I’m not doing anything with the forward() method other that what you can see here. I think it’s somehow related to how I’m using CIFAR, as I didn’t get the error just doing
x = torch.randn(1, 3, 32, 32)
out, activations = net(x)
Thanks!