diff --git a/README.md b/README.md
index f4b057e..e4510cd 100644
--- a/README.md
+++ b/README.md
@@ -133,6 +133,5 @@ Luxury food photograph of an italian Linguine pasta alle vongole dish with lots
 
 ### TODO
 
-- FLUX Dev implementation
 - LoRA adapters
 - Command line args
\ No newline at end of file
diff --git a/src/flux_1_schnell/flux.py b/src/flux_1_schnell/flux.py
index a034c68..e4da8b2 100644
--- a/src/flux_1_schnell/flux.py
+++ b/src/flux_1_schnell/flux.py
@@ -33,7 +33,7 @@ class Flux1:
     def generate_image(self, seed: int, prompt: str, config: Config = Config()) -> PIL.Image.Image:
         sigmas = get_sigmas(config.num_inference_steps)
         if self.is_dev:
-            sigmas = shift_sigmas(sigmas)
+            sigmas = shift_sigmas(sigmas, config.width, config.height)
         latents = LatentCreator.create(config.height, config.width, seed)
 
         t5_tokens = self.t5_tokenizer.tokenize(prompt)