Update pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

pages/2_😈_BlackBox_and_WhiteBox_Attacks.py

@@ -61,7 +61,7 @@ def run(model, tokenizer, embedidng_layer=None, _bar_text=None, bar=None, text='
 
         ori_output = model(**_input)
 
-        ori_output = ori_output['logits']
+        ori_output = ori_output['logits'][:, -1, :]
 
         # get noise
         ori_embedding = embedidng_layer(_input['input_ids']).detach()
@@ -87,7 +87,7 @@ def run(model, tokenizer, embedidng_layer=None, _bar_text=None, bar=None, text='
             _input_['inputs_embeds'] = perturbed_embedding
             outputs_perturbed = model(**_input_)
 
-            outputs_perturbed = outputs_perturbed['logits']
+            outputs_perturbed = outputs_perturbed['logits'][:, -1, :]
 
             loss = loss_funt(ori_output, outputs_perturbed)
             loss.backward()