Enable evaluating PFNs trained with pytorch/PFNs to be evaluated in Ax

Summary: This PR enables PFNs generally to work in our MAST evaluation, as the registry didn't quite work before, and it additionally allows to use training checkpoints from pytorch/PFNs to be used to do evaluations.

Reviewed By: Balandat

Differential Revision: D80944578

Author: SamuelGabriel

test_community/models/test_prior_fitted_network.py

@@ -179,9 +179,10 @@ def test_unpack_checkpoint(self):
 
         model = config.model.create_model()
 
+        state_dict = model.state_dict()
         checkpoint = {
             "config": config.to_dict(),
-            "model_state_dict": model.state_dict(),
+            "model_state_dict": state_dict,
         }
 
         loaded_model = PFNModel(
@@ -195,10 +196,10 @@ def test_unpack_checkpoint(self):
         loaded_state_dict = loaded_model.pfn.state_dict()
         self.assertEqual(
             sorted(loaded_state_dict.keys()),
-            sorted(model.state_dict().keys()),
+            sorted(state_dict.keys()),
         )
         for k in loaded_state_dict.keys():
-            self.assertTrue(torch.equal(loaded_state_dict[k], model.state_dict()[k]))
+            self.assertTrue(torch.equal(loaded_state_dict[k], state_dict[k]))
 
 
 class TestPriorFittedNetworkUtils(BotorchTestCase):