mend

mims-harvard · Nov 12, 2024 · f77297b · f77297b
1 parent 26355f3
commit f77297b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tdc/test/test_model_server.py b/tdc/test/test_model_server.py
@@ -154,7 +154,9 @@ def testGeneformerTokenizer(self):
                 # build an attention mask
                 attention_mask = torch.tensor(
                     [[x[0] != 0, x[1] != 0] for x in batch])
-                outputs = model(batch, attention_mask=attention_mask)
+                outputs = model(batch,
+                                attention_mask=attention_mask,
+                                output_hidden_states=True)
                 layer_to_quant = quant_layers(model) + (
                     -1
                 )  # TODO note this can be parametrized to either 0 (extract last embedding layer) or -1 (second-to-last which is more generalized)