fix: disable inference_mode entirely for training via inference_mode(False)

torch.enable_grad() alone is insufficient: operations on inference tensors
(created inside ComfyUI's outer inference_mode context) produce inference
tensors even inside enable_grad, breaking autograd. inference_mode(False)
exits the inference context so the deepcopy, apply_lora, and training loop
run with a fully clean autograd context.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-05 22:40:50 +02:00
parent 849f31e2a6
commit afb3242eca
+5 -4
View File
@@ -365,10 +365,11 @@ class SelvaLoraTrainer:
raise ValueError("[LoRA Trainer] No clips could be loaded.")
print(f"[LoRA Trainer] {len(dataset)} clip(s) ready.", flush=True)
# Everything from here runs inside enable_grad: ComfyUI wraps nodes in
# inference_mode, and nn.Parameters created in that context are inference
# tensors that can't enter autograd even with requires_grad=True.
with torch.enable_grad():
# ComfyUI executes nodes inside torch.inference_mode(). Inference tensors
# can't participate in autograd even with enable_grad — disable inference
# mode entirely so deepcopy, apply_lora, and the training loop all run
# with a clean autograd context.
with torch.inference_mode(False), torch.enable_grad():
return self._train_inner(
model, dataset, feature_utils_orig, seq_cfg,
device, dtype, variant, mode,