Hi, I am trying to visualize a pre-trained model’s intermediate activations using register_hooks and I got the error as mentioned in the above subject and as far as I know, my model and inputs are all on GPU itself. I cant understand why weights are loaded in the CPU as printing parameters show all params are on Cuda itself
Fixing seeds:
seed =0
random.seed(seed)
torch.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
np.random.seed(seed)

def seed_worker(worker_id):
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

Model:
class SimpleMLP(nn.Module):

    def __init__(self, num_classes=2, input_size=70, hidden_size=100):
        super(SimpleMLP, self).__init__()

        self.features = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(inplace=True),
            nn.Dropout(),
        )
        self.classifier = nn.Linear(hidden_size, num_classes)
        self._input_size = input_size

    def forward(self, x):
        x = x.contiguous()
        x = x.view(x.size(0), self._input_size)
        x = self.features(x)
        x = self.classifier(x)
        return x


Loading custom dataset:
class ToTensor:
    # Convert ndarrays to Tensors
    def __call__(self, sample):
        inputs, targets = sample
        inputs = np.array(inputs)
        targets = np.array([targets])
        return torch.from_numpy(inputs), torch.from_numpy(targets)

class IDSDataset(Dataset):
    def __init__(self):
        self.path_ds = './tSNEfigs/stratify/X_t.npy'
        self.path_lab = './tSNEfigs/stratify/y_t.npy'
        self.ds = np.load(self.path_ds) # each data point of shape (70,)
        self.label = np.load(self.path_lab) 
        self.n_samples = len(self.ds)
        self.transform = ToTensor()

    def __len__(self):
        return self.n_samples

    def __getitem__(self,index):
        sample = self.ds[index],int(self.label[index,-1])
        if self.transform:
            sample=self.transform(sample)
        return sample

Visualisation using tSNE:
test_ds=DataLoader(test_ds,batch_size=64,num_workers=4,worker_init_fn=seed_worker)

# Loading model
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model_task1 = SimpleMLP(num_classes=2,input_size=70,hidden_size=100)      # Model architecture
# Loading weights
model_task1.load_state_dict(torch.load('./weights/After_training_Task_1',map_location=device))
model_task1.to(device)
model_task1.features[2].register_forward_hook(get_activation('features[2]'))

print('Device params:\n ')
for name, param in model_task1.named_parameters():
    print(name, param.device)

out=[]
y_t=[]
model_task1.eval()  # Eval mode
with torch.no_grad():
    for batch_idx, (data, target) in enumerate(test_ds):
        data.to(device)     # Converting to cuda tensors
        target.to(device)   # Converting to cuda tensors
        y_t.append(target.cpu)
        output=model_task1(data)
        act = activation['features[2]']
        out.append(act)
    whole_x = [t.cpu.numpy() for t in out]
    whole_x = np.vstack(whole_x)
    whole_x = np.array(whole_x)

    tsne = TSNE(n_components=2, random_state=0,verbose=1,perplexity=50,init='pca')
    X_2d = tsne.fit_transform(whole_x)
    target_ids = [0,1] # (data_bin_hot['label'].nunique())
    y_t = y_t.astype('int')
    from matplotlib import pyplot as plt
    plt.figure(figsize=(10,6))
    colors = 'r', 'g'
    target_names=['Normal','Attack']
    for i, c, label in zip(target_ids, colors, target_names):
        plt.scatter(X_2d[y_t[:,-1] == i,0], X_2d[y_t[:,-1] == i,1], c=c, label=label)
    plt.legend()
    plt.savefig("./tSNEfigs/stratify/testinput.png")
    plt.show()