Upload lewm_model.py

lewm_model.py

@@ -302,6 +302,8 @@ class LeWorldModel(nn.Module):
         self.action_encoder = action_encoder
         self.projector = projector or nn.Identity()
         self.pred_proj = pred_proj or nn.Identity()
+        # SIGReg registered as a submodule so model.to(device) moves its buffers
+        self.sigreg = SIGReg()
 
     def encode(self, pixels: torch.Tensor) -> torch.Tensor:
         """
@@ -359,8 +361,8 @@ class LeWorldModel(nn.Module):
         pred_loss = (pred_emb[:, :-1] - emb[:, 1:history_size]).pow(2).mean()
 
         # SIGReg on step-wise embeddings (transpose to (T, B, D))
-        sigreg = SIGReg()
-        sigreg_loss = sigreg(emb.transpose(0, 1))
+        # self.sigreg is a registered submodule so it follows model.to(device)
+        sigreg_loss = self.sigreg(emb.transpose(0, 1))
 
         loss = pred_loss + sigreg_weight * sigreg_loss
         return {