Commit 41955edc authored by lizz's avatar lizz
Browse files

fix: dim shoud be int


Signed-off-by: default avatarlizz <lizz@sensetime.com>
parent 48024c15
......@@ -8,7 +8,8 @@ import torch
def maybe_dim_size(index, dim_size=None):
if dim_size is not None:
return dim_size
return index.max().item() + 1 if index.numel() > 0 else 0
dim = index.max().item() + 1 if index.numel() > 0 else 0
return int(dim)
def broadcast(src, index, dim):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment