So, back to decorating the classes and not the objects so that torch.save() does...