search

lix19937 / tensorrt-insight

Deep insight tensorrt, including but not limited to qat, ptq, plugin, triton_inference, cuda
12 stars 0 forks source link

torch.randint usage in msda plugin #48

Open lix19937 opened 1 month ago

lix19937 commented 1 month ago

https://pytorch.org/docs/stable/generated/torch.randint.html 


torch.randint(3, 5, (3,))

torch.randint(10, (2, 2))

torch.randint(3, 10, (2, 2))

>>> torch.randint(3, 5, (3,))
tensor([4, 3, 4])

>>> torch.randint(10, (2, 2))
tensor([[0, 2],
        [5, 5]])

>>> torch.randint(3, 10, (2, 2))
tensor([[4, 5],
        [6, 7]])