Update app.py

app.py

@@ -202,14 +202,16 @@ class MultiScaleRetention(nn.Module):
         
         outputs = []
         
+        # ✅ Decay를 입력과 같은 device로 이동
+        decay = torch.sigmoid(self.decay).view(1, -1, 1, 1).to(queries.device)
+        
         # Sequential processing (O(n))
         for t in range(seq_len):
             q_t = queries[:, :, t, :]  # [B, H, D]
             k_t = keys[:, :, t, :]     # [B, H, D]
             v_t = values[:, :, t, :]   # [B, H, D]
             
-            # Decay
-            decay = torch.sigmoid(self.decay).view(1, -1, 1, 1)
+            # Decay application
             state = decay * state
             
             # State update: S = decay * S + k @ v^T