make_grad_mat.py

        sim_mat_dims = (len(dl.dataset), len(dl.dataset))
        print("Dimensions of similarity matrix is", sim_mat_dims)
        print("Making empty matrix to store similarities ......")
        feat_mat = np.empty(sim_mat_dims, dtype=np.float32)
        loss_fn = nn.CrossEntropyLoss(reduction='mean').to(self.device)
        for idx, data in tqdm(enumerate(dl)):
            loss_val = loss_fn(net(data[0].to(self.device)), data[1].to(self.device))
            grad_list = torch.autograd.grad(loss_val, inputs = [p for p in net.parameters() if p.requires_grad])
            feats_outer = [t.flatten() for t in grad_list]
            feats_outer = torch.cat(feats_outer)
            for idx, data in tqdm(enumerate(dl)):
                loss_val = loss_fn(net(data[0].to(self.device)), data[1].to(self.device))
                grad_list = torch.autograd.grad(loss_val, inputs = [p for p in net.parameters() if p.requires_grad])
                feats_inner = [t.flatten() for t in grad_list]
                feats_inner = torch.cat(feats_inner)
                feat_mat[idx, idx] = torch.norm(feats_outer - feats_inner).cpu().detach().numpy()

The first thing to note about that code is that there is a bug, since you are using the same variable names in the inner loop as you are in the outer loop. This is scope masking and also the final feat_mat[idx,idx] = is only setting the diagonal of the matrix. Updated version below but in general in your code you should always use distinct variable names in nested loops:

        sim_mat_dims = (len(dl.dataset), len(dl.dataset))
        print("Dimensions of similarity matrix is", sim_mat_dims)
        print("Making empty matrix to store similarities ......")
        feat_mat = np.empty(sim_mat_dims, dtype=np.float32)
        loss_fn = nn.CrossEntropyLoss(reduction='mean').to(self.device)
        for idx1, data1 in tqdm(enumerate(dl)):
            loss_val1 = loss_fn(net(data1[0].to(self.device)), data1[1].to(self.device))
            grad_list1 = torch.autograd.grad(loss_val1, inputs1 = [p for p in net.parameters() if p.requires_grad])
            feats_outer = [t.flatten() for t in grad_list1]
            feats_outer = torch.cat(feats_outer)
            for idx2, data2 in tqdm(enumerate(dl)):
                loss_val2 = loss_fn(net(data2[0].to(self.device)), data2[1].to(self.device))
                grad_list2 = torch.autograd.grad(loss_val2, inputs2 = [p for p in net.parameters() if p.requires_grad])
                feats_inner = [t.flatten() for t in grad_list2]
                feats_inner = torch.cat(feats_inner)
                feat_mat[idx1, idx2] = torch.norm(feats_outer - feats_inner).cpu().detach().numpy()