Enhance Infinity model loading by specifying device type in autocast context for improved performance

app.py

@@ -198,7 +198,7 @@ def load_infinity(
     print(f'[Loading Infinity]')
     text_maxlen = 512
     torch.cuda.empty_cache()
-    with torch.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True), torch.no_grad():
+    with torch.amp.autocast(device_type=device, enabled=True, dtype=torch.bfloat16, cache_enabled=True), torch.no_grad():
         infinity_test: Infinity = Infinity(
             vae_local=vae, text_channels=text_channels, text_maxlen=text_maxlen,
             shared_aln=True, raw_scale_schedule=scale_schedule,