I found it… One of the scripts from a different part of my pipeline that was imported but wasn’t being used at all had a torch.no_grad in it facepalm
In any case thanks for all the help and sorry for the wild goose chase