Upload SmalLmForCausalLM

model.py

@@ -635,6 +635,8 @@ class SmalLmModel(SmalLmPreTrainedModel):
         cache_position: Optional[torch.Tensor],
     ):
         if USE_FLASH and inputs_embeds.is_cuda:
+            if attention_mask is None:
+                attention_mask = torch.ones(*inputs_embeds.shape[:2], device=inputs_embeds.device)
             return attention_mask
         dtype, device = inputs_embeds.dtype, inputs_embeds.device
         past_token = (