From f77297b5eb05ec3a0b8c9c214354f36dda5fe796 Mon Sep 17 00:00:00 2001 From: Alejandro Velez-Arce Date: Mon, 11 Nov 2024 19:53:52 -0500 Subject: [PATCH] mend --- tdc/test/test_model_server.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tdc/test/test_model_server.py b/tdc/test/test_model_server.py index 29a1df09..6a51ab5c 100644 --- 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)