You’re right, it eventually converges, but I don’t understand how the gradients for a single sample can push it in the wrong direction for multiple steps.
My basic understanding of SGD says that with a single sample, the gradients should be guaranteed to be in the right direction. So how is this happening?
Here is my model:
class Net(torch.nn.Module):
    def __init__(self, num_node_features):
        super().__init__()
        self.num_node_features = num_node_features
        self.alpha = torch.nn.Parameter(torch.randn(1))
        self.linear_1 = torch.nn.Linear(num_node_features, num_node_features * 2)
        self.linear_2 = torch.nn.Linear(self.linear_1.out_features, self.linear_1.out_features)
        self.linear_3 = torch.nn.Linear(self.linear_2.out_features, 1)
        self.linear_4 = torch.nn.Linear(self.linear_2.out_features * 2, self.linear_2.out_features)
        self.linear_5 = torch.nn.Linear(self.linear_2.out_features, self.linear_1.in_features)
        self.linear_6 = torch.nn.Linear(self.linear_1.in_features, 1)

    def forward(self, protein_ligand_pair):
        graph_embeddings = []

        for molecule in protein_ligand_pair:
            #These are actually batches of matrices.
            node_matrix, adjacency_matrix = molecule
            propagations = 5
            for step in range(propagations):
                smoothed_node_matrix = torch.matmul(adjacency_matrix, node_matrix)
                node_matrix = self.alpha*node_matrix + (1-self.alpha)*smoothed_node_matrix

                batch_size, num_nodes, num_input_features = node_matrix.shape
                new_node_matrix = torch.empty(batch_size, num_nodes, self.linear_2.in_features, device = self.alpha.device)
                for node_i in range(num_nodes):
                    linear_layer = self.linear_1 if step == 0 else self.linear_2
                    node_features = node_matrix[:, node_i]
                    new_node_matrix[:, node_i] = torch.nn.ReLU()(linear_layer(node_features))
                node_matrix = new_node_matrix
            
            num_nodes, num_features = node_matrix.shape[1:]
            aggregation_weights = torch.empty(batch_size, num_nodes, device = self.alpha.device)
            for node_i in range(num_nodes):
                node_features = node_matrix[:, node_i]
                aggregation_weights[:, node_i] = torch.nn.ReLU()(self.linear_3(node_features))

            aggregate = torch.matmul(aggregation_weights, node_matrix)
            graph_embeddings.append(aggregate)
        
        protein_ligand_concatenated = torch.cat(graph_embeddings, axis = 2)
        last_hidden_output = torch.nn.Sequential(self.linear_4, torch.nn.ReLU(), self.linear_5, torch.nn.ReLU())(protein_ligand_concatenated)
        ligand_protein_affinity = torch.nn.Sigmoid()(self.linear_6(last_hidden_output)).flatten()

        return ligand_protein_affinity

Some context for the “for step in range(propagations)” loop: How does applying the same convolutional layer to its own output affect learning?