[FRONTEND] Fix pytorch warning. (#560)

UserWarning: __floordiv__ is deprecated, and its behavior will change in a future version of pytorch. It currently rounds toward 0 (like the 'trunc' function NOT 'floor'). This results in incorrect rounding for negative values. To keep the current behavior, use torch.div(a, b, rounding_mode='trunc').
This commit is contained in:
Keren Zhou
2022-06-20 20:12:09 -07:00
committed by GitHub
parent b5e728cb14
commit a428cf0bb2

View File

@@ -5,7 +5,7 @@ import triton
def rounded_linspace(low, high, steps, div):
ret = torch.linspace(low, high, steps)
ret = (ret.int() + div - 1) // div * div
ret = torch.div(ret.int() + div - 1, div, rounding_mode='trunc') * div
ret = torch.unique(ret)
return list(map(int, ret))