Add comment with pytorch issue
@@ -27,6 +27,7 @@ class Distance(torch.nn.Module):
and x1.dim() < 3
and x1.size(-2) <= 512
and x2.size(-2) <= 512
+ # TODO: remove when https://github.com/pytorch/pytorch/issues/22353 is fixed
and not (x1.is_cuda or x2.is_cuda)
):
# torch cdist doesn't support batch, also only available after pytorch 1.1
@@ -58,6 +59,7 @@ class Distance(torch.nn.Module):
and x1.dim() < 3
and x1.size(-2) <= 512
and x2.size(-2) <= 512
+ # TODO: remove when https://github.com/pytorch/pytorch/issues/22353 is fixed
and not (x1.is_cuda or x2.is_cuda)
):
# torch cdist doesn't support batch, also only available after pytorch 1.1
