When I first pass alll the nodes to cpu with the following code:
    def to_cpu(self):
        topological_order = nx.algorithms.dag.topological_sort(self)

        for node in topological_order:
            # try:
            node_info = self.nodes[node]
            node_info['op'].cpu()

When I call:
self.graph.to_cpu()
self.optimizer = torch.optim.Adam(self.graph.parameters())

I get the following error message:
UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the gradient for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations.