ultralytics 8.1.9 replace .size(0) with .shape[0] (#7957)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-02-01 23:33:02 +01:00 committed by GitHub
parent 66b32bb4dd
commit 70d4a3752e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 42 additions and 22 deletions

View file

@ -121,7 +121,7 @@ class MaskDecoder(nn.Module):
"""
# Concatenate output tokens
output_tokens = torch.cat([self.iou_token.weight, self.mask_tokens.weight], dim=0)
output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.size(0), -1, -1)
output_tokens = output_tokens.unsqueeze(0).expand(sparse_prompt_embeddings.shape[0], -1, -1)
tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1)
# Expand per-image data in batch direction to be per-mask