diff --git a/stable_audio_tools/inference/generation.py b/stable_audio_tools/inference/generation.py
index 9944dd86..4e81871f 100644
--- a/stable_audio_tools/inference/generation.py
+++ b/stable_audio_tools/inference/generation.py
@@ -91,7 +91,7 @@ def generate_diffusion_uncond(
 def generate_diffusion_cond(
         model,
         steps: int = 250,
-        cfg_scale=6,
+        cfg_scale: float = 6,
         conditioning: dict = None,
         conditioning_tensors: tp.Optional[dict] = None,
         negative_conditioning: dict = None,
@@ -427,4 +427,4 @@ def build_mask(sample_size, mask_args):
     if marination > 0:        
         mask = mask * (1-marination) 
     #print(mask)
-    return mask
\ No newline at end of file
+    return mask