Debug network_swin2sr

mv-lab · Sep 27, 2022 · 0abb1e0 · 0abb1e0
1 parent fd9c1ad
commit 0abb1e0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/models/network_swin2sr.py b/models/network_swin2sr.py
@@ -148,7 +148,7 @@ def forward(self, x, mask=None):
 
  # cosine attention
  attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
- logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01))).exp()
+ logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01)).to(self.logit_scale.device)).exp()
  attn = attn * logit_scale
 
  relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)