change cpu

inference.py

@@ -12,6 +12,8 @@ from utils.image_io import save_image_tensor
 
 from text_net.model import AirNet
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 def test_Derain_Dehaze(opt, net, dataset, task="derain"):
     output_path = opt.output_path + task + '/'
     subprocess.check_output(['mkdir', '-p', output_path])
@@ -39,7 +41,6 @@ def infer(text_prompt = "", img=None):
 
     opt = parser.parse_args()
     # opt.text_prompt = text_prompt
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
     np.random.seed(0)
     torch.manual_seed(0)