Issue with half precision training
Update swin_unetr.py
fixed potential issue with discrepant data types
@@ -503,7 +503,7 @@ class WindowAttention(nn.Module):
else:
attn = self.softmax(attn)
- attn = self.attn_drop(attn)
+ attn = self.attn_drop(attn).to(v.dtype)
x = (attn @ v).transpose(1, 2).reshape(b, n, c)
x = self.proj(x)
x = self.proj_drop(x)
