It also looks like subsetting the features in x is causing this problem. If I don’t do that, the gradients appear to be able to flow. This step appears to be causing the problem and I have no idea why.
        # Take the d p vectors, and concat them, followed by putting classifier on top 
        stop_num = int(edge_index.shape[1]/2)
        
        node_idx = edge_index.T[0:stop_num]
        
        d_rep = x[node_idx[:,0]]
        p_rep = x[node_idx[:,1]]
        edge_weight_half = edge_weight[0:stop_num].type(torch.FloatTensor).view(-1,1).cuda()
        
        d_p_pair = torch.cat((d_rep, p_rep, edge_weight_half), dim=1)