Update model.py

model.py

@@ -54,19 +54,17 @@ class ArcMarginProduct(nn.Module):
 class DualStreamTransformer(nn.Module):
     def __init__(self, feat_num_1, feat_num_2, d_model=64, num_classes=3, dropout=0.3):
         super().__init__()
-        # Stream 1: CCMQ
         self.feat_tokenizers_1 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(feat_num_1)])
         enc_layer_1 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=128, dropout=dropout, batch_first=True)
         self.encoder_1 = nn.TransformerEncoder(enc_layer_1, num_layers=2)
         self.cls_token_1 = nn.Parameter(torch.zeros(1, 1, d_model))
         
-        # Stream 2: OSDI
         self.feat_tokenizers_2 = nn.ModuleList([nn.Linear(1, d_model) for _ in range(feat_num_2)])
         enc_layer_2 = nn.TransformerEncoderLayer(d_model=d_model, nhead=4, dim_feedforward=128, dropout=dropout, batch_first=True)
         self.encoder_2 = nn.TransformerEncoder(enc_layer_2, num_layers=2)
         self.cls_token_2 = nn.Parameter(torch.zeros(1, 1, d_model))
 
-        # Fusion
+
         self.fusion = nn.Sequential(
             nn.Linear(d_model * 2, d_model),
             nn.LayerNorm(d_model),
@@ -89,5 +87,5 @@ class DualStreamTransformer(nn.Module):
     def forward(self, x1, x2):
         feat_1 = self.forward_stream(x1, self.feat_tokenizers_1, self.encoder_1, self.cls_token_1)
         feat_2 = self.forward_stream(x2, self.feat_tokenizers_2, self.encoder_2, self.cls_token_2)
-        combined = torch.cat((feat_1, feat2), dim=1)
+        combined = torch.cat((feat_1, feat_2), dim=1)
         return self.fusion(combined)