Dima fix imitation learning (mila-iqia#93)

* make sure instr_embedding and mask have the same length

* fix typos in the comments

babyai/model.py

@@ -211,11 +211,17 @@ def forward(self, obs, memory, instr_embedding=None):
             # outputs: B x L x D
             # memory: B x M
             mask = (obs.instr != 0).float()
-
-            # The mask tensor should be always at least a big as instr_embedding.
-            # It can be too big though, because instr_embeddings might be shorter than obs.instr
-            # when the last position in a batch is all zeros.
+            # The mask tensor has the same length as obs.instr, and
+            # thus can be both shorter and longer than instr_embedding.
+            # It can be longer if instr_embedding is computed
+            # for a subbatch of obs.instr.
+            # It can be shorter if obs.instr is a subbatch of
+            # the batch that instr_embeddings was computed for.
+            # Here, we make sure that mask and instr_embeddings
+            # have equal length along dimension 1.
             mask = mask[:, :instr_embedding.shape[1]]
+            instr_embedding = instr_embedding[:, :mask.shape[1]]
+
             keys = self.memory2key(memory)
             pre_softmax = (keys[:, None, :] * instr_embedding).sum(2) + 1000 * mask
             attention = F.softmax(pre_softmax, dim=1)