Update lwm_model.py

lwm_model.py

@@ -2,21 +2,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-import random
-
-# Set manual seed for reproducibility
-def set_random_seed(seed=42):
-    torch.manual_seed(seed)
-    np.random.seed(seed)
-    random.seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-    # Ensures deterministic behavior
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
-
-# Call the seed function
-set_random_seed()
 
 ELEMENT_LENGTH = 16
 D_MODEL = 64
@@ -53,7 +38,7 @@ class Embedding(nn.Module):
         seq_len = x.size(1)
         pos = torch.arange(seq_len, dtype=torch.long, device=x.device)
         pos = pos.unsqueeze(0).expand_as(x[:, :, 0])
-        tok_emb = self.proj(x.float())  # Ensure consistency in floating-point precision
+        tok_emb = self.proj(x.float())  
         embedding = tok_emb + self.pos_embed(pos)
         return self.norm(embedding)
 
@@ -124,15 +109,12 @@ class LWM(torch.nn.Module):
         embed_weight = self.embedding.proj.weight
         d_model, n_dim = embed_weight.size()
         self.decoder = nn.Linear(d_model, n_dim, bias=False)
-        self.decoder.weight = nn.Parameter(embed_weight.transpose(0, 1))
         self.decoder_bias = nn.Parameter(torch.zeros(n_dim))
 
     @classmethod
     def from_pretrained(cls, ckpt_name='model_weights.pth', device='cuda', use_auth_token=None):
-        # Define model
         model = cls().to(device)
 
-        # Load model weights
         ckpt_path = ckpt_name
         model.load_state_dict(torch.load(ckpt_path, map_location=device))
         print(f"Model loaded successfully from {ckpt_path} to {device}")