loaded with bf16
This commit is contained in:
parent
33a3626b74
commit
b89670ad6c
|
@ -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")
|
||||
|
|
Loading…
Reference in New Issue