use better context manager to fix potential problems

lllyasviel · Feb 8, 2024 · 291ec74 · 291ec74
1 parent 760f727
commit 291ec74
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules_forge/forge_loader.py b/modules_forge/forge_loader.py
@@ -130,7 +130,7 @@ class WeightsLoader(torch.nn.Module):
  return ForgeSD(model_patcher, clip, vae, clipvision)
 
 
-@torch.inference_mode()
+@torch.no_grad()
 def load_model_for_a1111(timer, checkpoint_info=None, state_dict=None):
  a1111_config_filename = find_checkpoint_config(state_dict, checkpoint_info)
  a1111_config = OmegaConf.load(a1111_config_filename)