Update iris/src/tokenizer.py

iris/src/tokenizer.py

@@ -69,7 +69,7 @@ class Tokenizer(nn.Module):
         return rec
 
     def decode_obs_tokens(self, obs_tokens, num_observations_tokens):
-        embedded_tokens = self.embedding(self.obs_tokens)     # (B, K, E)
+        embedded_tokens = self.embedding(obs_tokens)     # (B, K, E)
         z = rearrange(embedded_tokens, 'b (h w) e -> b e h w', h=int(np.sqrt(num_observations_tokens)))
         rec = self.decode(z, should_postprocess=True)         # (B, C, H, W)
         return torch.clamp(rec, 0, 1)