fix eos token for batched case

h_transformer_1d/autoregressive_wrapper.py

@@ -4,6 +4,9 @@
 
 # helper function
 
+def exists(val):
+    return val is not None
+
 def eval_decorator(fn):
     def inner(model, *args, **kwargs):
         was_training = model.training
@@ -56,8 +59,15 @@ def generate(self, start_tokens, seq_len, eos_token = None, temperature = 1., fi
 
             out = torch.cat((out, sample), dim=-1)
 
-            if eos_token is not None and (sample == eos_token).all():
-                break
+            if exists(eos_token):
+                is_eos_token = (out == eos_token)
+
+                if is_eos_token.any(dim = -1).all():
+                    # mask out everything after the eos tokens
+                    shifted_is_eos_tokens = F.pad(is_eos_tokens, (1, -1))
+                    mask = shifted_is_eos_tokens.float().cumsum(dim = -1) >= 1
+                    out = out.masked_fill(mask, self.pad_value)
+                    break
 
         out = out[:, t:]
 

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'h-transformer-1d',
   packages = find_packages(),
-  version = '0.1.0',
+  version = '0.1.1',
   license='MIT',
   description = 'H-Transformer 1D - Pytorch',
   author = 'Phil Wang',