Generate a sample input tensor from a shape string like '1,2048'.
(
input_shape: str,
dtype: torch.dtype,
device: str = "cuda",
seed: int = 42,
)
| 200 | # --------------------------------------------------------------------------- |
| 201 | |
| 202 | def generate_sample_input( |
| 203 | input_shape: str, |
| 204 | dtype: torch.dtype, |
| 205 | device: str = "cuda", |
| 206 | seed: int = 42, |
| 207 | ) -> torch.Tensor: |
| 208 | """Generate a sample input tensor from a shape string like '1,2048'.""" |
| 209 | dims = [int(d.strip()) for d in input_shape.split(",")] |
| 210 | torch.manual_seed(seed) |
| 211 | |
| 212 | if dtype in (torch.int32, torch.int64, torch.long): |
| 213 | # For language models, generate token IDs (assume vocab size ~32000) |
| 214 | return torch.randint(0, 32000, dims, device=device, dtype=dtype) |
| 215 | else: |
| 216 | return torch.randn(dims, device=device, dtype=dtype) |
| 217 | |
| 218 | |
| 219 | def infer_input_type(model: nn.Module) -> str: |