diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py
index c90a48d..f655fcf 100644
--- a/generator/ctrl/ctrl_generator.py
+++ b/generator/ctrl/ctrl_generator.py
@@ -317,7 +317,7 @@ class CTRLGenerator():
 
             import pdb
             pdb.set_trace()
-            tokens_generated_so_far = ' '.join([self.idx2word[c] for c in tokens_generated[0][len(text):token+2].squeeze()])
+            tokens_generated_so_far = ' '.join([self.idx2word[c] for c in tokens_generated[0][len(text):token+2]])
             tokens_generated_so_far = re.sub('(@@ )', '', string=tokens_generated_so_far)
             tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
             result = tokens_generated_so_far