Solved by using out-of-place operation for the following lines of code:
                        graph.cells[c].nodes[n].connections[cc].combined_feature_map = \
                            graph.cells[c].nodes[n].connections[cc].combined_feature_map + \
                            graph.cells[c].nodes[n].connections[cc].edges[e].forward_f(x)  # Ltrain(w±, alpha)

Note: I turned on torch.autograd.set_detect_anomaly(True) to have found this solution.