From b89670ad6c1089c4fda8d7332a20fa8f34956a58 Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Mon, 24 Feb 2025 15:21:08 +0100 Subject: [PATCH] loaded with bf16 --- src/aiunn/finetune.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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")