From 0fa14c4a6f120a94f7fdf6846a210c4fa1115e1f Mon Sep 17 00:00:00 2001 From: Falko Habel Date: Mon, 24 Feb 2025 15:22:41 +0100 Subject: [PATCH] added gpu support --- 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 19a2aaa..ec18832 100644 --- a/src/aiunn/finetune.py +++ b/src/aiunn/finetune.py @@ -97,7 +97,7 @@ pretrained_model_path = "/root/vision/AIIA/AIIA-base-512" 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") +device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # Move model to device using channels_last memory format. model = model.to(device, memory_format=torch.channels_last)