compile is back

scripts/exp/train.py

@@ -513,14 +513,14 @@ def load(
 
     if args["fine_tune"]:
         assert fine_tune_checkpoint is not None, "Must provide a fine-tune checkpoint"
-        model = (
+        model = torch.compile(
             VampNet.load(location=Path(fine_tune_checkpoint), 
                          map_location="cpu", 
             )
         )
 
 
-    model = VampNet() if model is None else model
+    model = torch.compile(VampNet()) if model is None else model
     model = accel.prepare_model(model)
 
     # assert accel.unwrap(model).n_codebooks == codec.quantizer.n_codebooks