try fix #1761 again (#1680)

2026-08-16 13:13:16 +02:00 · 2023-12-31 04:38:31 -08:00
parent 2521af308b
commit 0c4f20a0d2
2 changed files with 3 additions and 4 deletions
@@ -108,8 +108,7 @@ class Resampler(nn.Module):
            )

    def forward(self, x):
-        
-        latents = self.latents.repeat(x.size(0), 1, 1)
+        latents = self.latents.repeat(x.size(0), 1, 1).to(x)
        
        x = self.proj_in(x)
        
@@ -118,4 +117,4 @@ class Resampler(nn.Module):
            latents = ff(latents) + latents
            
        latents = self.proj_out(latents)
-        return self.norm_out(latents)
+        return self.norm_out(latents)