Update modeling_alpha3d.py

modeling_alpha3d.py

@@ -10,19 +10,26 @@ class Alpha3DModel(PreTrainedModel):
         super().__init__(config)
         self.config = config
         
-        layers = []
-        input_dim = 5
-        
-        prev_dim = input_dim
-        for h_dim in config.hidden_dims:
-            layers.append(nn.Linear(prev_dim, h_dim))
-            layers.append(nn.BatchNorm1d(h_dim))
-            layers.append(nn.ReLU())
-            prev_dim = h_dim
+        # ТОЧНАЯ КОПИЯ ТОГО, ЧТО БЫЛО ПРИ ОБУЧЕНИИ
+        # Без циклов, чтобы индексы (0, 1, ... 8) совпали идеально
+        self.net = nn.Sequential(
+            # Слой 1
+            nn.Linear(5, 128),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
             
-        layers.append(nn.Linear(prev_dim, config.num_points * 6))
-        
-        self.net = nn.Sequential(*layers)
+            # Слой 2
+            nn.Linear(128, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            
+            # Слой 3 (Тут НЕ БЫЛО BatchNorm при обучении!)
+            nn.Linear(512, 1024),
+            nn.ReLU(),
+            
+            # Выход
+            nn.Linear(1024, config.num_points * 6)
+        )
 
     def forward(self, x):
         out = self.net(x)