finetune_class #1
|
@ -94,7 +94,7 @@ transform = transforms.Compose([
|
||||||
|
|
||||||
# Load the base AIIABase model and wrap it with the Upsampler.
|
# Load the base AIIABase model and wrap it with the Upsampler.
|
||||||
pretrained_model_path = "/root/vision/AIIA/AIIA-base-512"
|
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)
|
model = Upsampler(base_model)
|
||||||
|
|
||||||
device = torch.device("cpu")#torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
device = torch.device("cpu")#torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||||
|
|
Loading…
Reference in New Issue