Update modeling_quiet.py

modeling_quiet.py

@@ -1072,16 +1072,7 @@ class QuietModel(QuietPreTrainedModel):
                 )
 
         if attention_mask is None:
-            if input_ids is not None:
-                attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
-            elif inputs_embeds is not None:
-                attention_mask = torch.ones(
-                    (batch_size, seq_len),
-                    dtype=torch.bool,
-                    device=inputs_embeds.device
-                )
-            else:
-                raise ValueError("Either input_ids or inputs_embeds should be provided.")
+            attention_mask = torch.ones((batch_size, seq_len), dtype=torch.bool, device=inputs_embeds.device)
 
         if attention_mask.dim() == 2:
             attention_mask = attention_mask.view(batch_size, 1, 1, seq_len)
@@ -1091,7 +1082,7 @@ class QuietModel(QuietPreTrainedModel):
         elif attention_mask.dim() != 4:
             raise ValueError(f"Attention mask should be of shape (batch_size, 1, seq_len, seq_len) or (batch_size, 1, 1, seq_len), but got {attention_mask.shape}")
 
-        attention_mask = attention_mask.to(dtype=torch.bool, device=input_ids.device)
+        attention_mask = attention_mask.to(dtype=torch.bool)
 
         hidden_states = inputs_embeds