albanD

import torch
from torch.autograd import Function
import torch.utils._pytree as pytree

# Basically wraps things in and out before passing it to the real function that the user defined.
def pytreeify(cls):
    assert issubclass(cls, Function)

    orig_fw = cls.forward
    orig_bw = cls.backward