Make the class implementations compatible with Torch compile #178
Description
Ensure that all the code works as expected with torch.compile. It would be good to include this in the unit tests as well.
{{ message }}
Ensure that all the code works as expected with torch.compile. It would be good to include this in the unit tests as well.