Hotfix

bert.py

@@ -59,7 +59,7 @@ class BertSelfAttention(nn.Module):
 
   def forward(self, hidden_states, attention_mask):
     """
-    hidden_states: [bs, seq_len, hidden_state]
+    hidden_states: [bs, seq_len, hidden_size]
     attention_mask: [bs, 1, 1, seq_len]
     output: [bs, seq_len, hidden_state]
     """

constants.py

@@ -28,8 +28,8 @@ STSB_DEV = os.path.join(DATA_DIR, 'stsb-dev.parquet')
 # Training-specific constants
 SEED=11711
 NUM_CPU_CORES=4
-EPOCHS=1
-USE_GPU=False
+EPOCHS=10
+USE_GPU=True
 BATCH_SIZE_CSE=8
 BATCH_SIZE_SST=64
 BATCH_SIZE_CFIMDB=8