diff --git a/src/aiunn/finetune.py b/src/aiunn/finetune.py
index 80088bb..19a2aaa 100644
--- a/src/aiunn/finetune.py
+++ b/src/aiunn/finetune.py
@@ -94,7 +94,7 @@ transform = transforms.Compose([
 
 # Load the base AIIABase model and wrap it with the Upsampler.
 pretrained_model_path = "/root/vision/AIIA/AIIA-base-512"
-base_model = AIIABase.load(pretrained_model_path)
+base_model = AIIABase.load(pretrained_model_path, precision="bf16")
 model = Upsampler(base_model)
 
 device = torch.device("cpu")#torch.device("cuda" if torch.cuda.is_available() else "cpu")