Hi,
As @iffiX mentioned, we do not compile graph in PyTorch, it will be constructed during data flow in forward pass. So, for an optimizer, you just define a class that accepts parameters of a model = nn.Module and implements step function for it. Literally, based on your code, you need to only remove last line self._train_fn = ... as there is no compile stage in PyTorch.
Here is the SGD implementation:
https://pytorch.org/docs/stable/_modules/torch/optim/sgd.html#SGD 3
Bests