loaded with bf16

This commit is contained in:
Falko Victor Habel 2025-02-24 15:21:08 +01:00
parent 33a3626b74
commit b89670ad6c
1 changed files with 1 additions and 1 deletions

View File

@ -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")