diff --git a/README.md b/README.md
index 5061bd1547..a6ba944a9e 100644
--- a/README.md
+++ b/README.md
@@ -34,6 +34,7 @@ import numpy as np
 
 generator = torch.Generator()
 generator = generator.manual_seed(6694729458485568)
+torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 1. Load models
 scheduler = GaussianDDPMScheduler.from_config("fusing/ddpm-lsun-church")