fix: torch.compile

train.py

@@ -142,6 +142,9 @@ elif args.model == "resnet101":
 else:
     raise NotImplementedError()
 
+if torch.__version__ >= "2.0":
+    model = torch.compile(model)
+
 detector = FontDetector(
     model=model,
     lambda_font=lambda_font,
@@ -154,8 +157,5 @@ detector = FontDetector(
     num_epochs=num_epochs,
 )
 
-if torch.__version__ >= "2.0":
-    detector = torch.compile(detector)
-
 trainer.fit(detector, datamodule=data_module, ckpt_path=args.checkpoint)
 trainer.test(detector, datamodule=data_module)