Skip to content

Commit 07c9c2d

Browse files
committed
chore: clean up for random int type input
1 parent 89d4a98 commit 07c9c2d

File tree

1 file changed

+3
-10
lines changed

1 file changed

+3
-10
lines changed

py/torch_tensorrt/_Input.py

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -364,18 +364,11 @@ def example_tensor(
364364
)
365365

366366
if isinstance(self.shape, dict):
367-
if (
368-
self.dtype == dtype.u8
369-
or self.dtype == dtype.i8
370-
or self.dtype == dtype.i32
371-
or self.dtype == dtype.i64
372-
):
367+
if self.dtype in [dtype.u8, dtype.i8, dtype.i32, dtype.i64]:
373368
type = self.dtype.to(torch.dtype, use_default=True)
374-
min_value = torch.iinfo(type).min
375-
max_value = torch.iinfo(type).max
376369
return torch.randint(
377-
min_value,
378-
max_value,
370+
torch.iinfo(type).min,
371+
torch.iinfo(type).max,
379372
self.shape[optimization_profile_field],
380373
dtype=type,
381374
)

0 commit comments

Comments
 (0)