diff --git a/test.py b/test.py
index 69cd9de..e52d1a5 100644
--- a/test.py
+++ b/test.py
@@ -282,7 +282,7 @@ def get_id(img_path):
 if torch.cuda.get_device_capability()[0]>6: # should be >=7
     print("Compiling model...")
     # https://huggingface.co/docs/diffusers/main/en/optimization/torch2.0
-    model_structure = torch.compile(model_structure, mode="reduce-overhead", fullgraph=True) # pytorch 2.0
+    model_structure = torch.compile(model_structure, mode="default", dynamic=True) # pytorch 2.0
 
 model = load_network(model_structure)
 
diff --git a/train.py b/train.py
index 1766bf5..dfffa0b 100755
--- a/train.py
+++ b/train.py
@@ -514,7 +514,7 @@ def save_network(network, epoch_label):
 if torch.cuda.get_device_capability()[0]>6: # should be >=7
     print("Compiling model... The first epoch may be slow, which is expected!")
     # https://huggingface.co/docs/diffusers/main/en/optimization/torch2.0
-    model = torch.compile(model, mode="reduce-overhead", fullgraph=True) # pytorch 2.0
+    model = torch.compile(model, mode="reduce-overhead", dynamic = True) # pytorch 2.0
 
 optim_name = optim.SGD #apex.optimizers.FusedSGD
 if opt.FSGD: # apex is needed