Generate appropriate sample input for the model.
(
model: nn.Module,
input_shape: List[int],
dtype: torch.dtype,
device: str,
)
| 341 | |
| 342 | |
| 343 | def generate_input( |
| 344 | model: nn.Module, |
| 345 | input_shape: List[int], |
| 346 | dtype: torch.dtype, |
| 347 | device: str, |
| 348 | ) -> Dict[str, Any]: |
| 349 | """Generate appropriate sample input for the model.""" |
| 350 | if _is_language_model(model): |
| 351 | # Language model: generate integer token IDs |
| 352 | batch = input_shape[0] if len(input_shape) >= 1 else 1 |
| 353 | seq_len = input_shape[1] if len(input_shape) >= 2 else 512 |
| 354 | input_ids = torch.randint(0, 32000, (batch, seq_len), device=device, dtype=torch.long) |
| 355 | return {"input_ids": input_ids} |
| 356 | else: |
| 357 | # Generic model: generate float tensor of given shape |
| 358 | x = torch.randn(*input_shape, device=device, dtype=dtype) |
| 359 | return {"x": x} |
| 360 | |
| 361 | |
| 362 | def _try_forward( |
no test coverage detected