fix: torch.set_grad_enabled(False) is a no-op outside a with block

[andrewwhitecdw]

This PR corrects the documentation in README.md: torch.set_grad_enabled(False) is a no-op outside a with block.

Changes

• README.md: torch.set_grad_enabled(False) is a no-op outside a with block.

Details

--- a/README.md
+++ b/README.md
@@ -1,10 +1,10 @@
-torch.set_grad_enabled(False)
-pipe = CausalDiffusionInferencePipeline(config, device=device)
-load_generator_checkpoint(pipe.generator, merged_checkpoint_path)
-pipe = pipe.to(device=device, dtype=torch.bfloat16)
-place_vae_for_streaming(pipe, config)  # honor streaming_vae + vae_device when set
-pipe.generator.model.eval().requires_grad_(False)
-
-noise, prompts = prepare_single_prompt_inputs(config, prompt, device)
-video = pipe.inference(noise=noise, text_prompts=prompts)
-save_video(video[0], "videos/quickstart/sample.mp4", fps=24)
+with torch.inference_mode():
+    pipe = CausalDiffusionInferencePipeline(config, device=device)
+    load_generator_checkpoint(pipe.generator, merged_checkpoint_path)
+    pipe = pipe.to(device=device, dtype=torch.bfloat16)
+    place_vae_for_streaming(pipe, config)  # honor streaming_vae + vae_device when set
+    pipe.generator.model.eval().requires_grad_(False)
+
+    noise, prompts = prepare_single_prompt_inputs(config, prompt, device)
+    video = pipe.inference(noise=noise, text_prompts=prompts)
+    save_video(video[0], "videos/quickstart/sample.mp4", fps=24)

Tests

Let me know if you want tests added for this fix or not.