Potential bug in the transformer.py and detr_vae.py?

[Facebear-ljx]

Hi, thanks for the amazing work! Recently, I tried to reimplement your code. However, I found something strange.

Especially, the forward of Transformer class in transformer.py only returns hs，which is a tensor with shape [Decoder block num, B, num_query, C]:

act/detr/models/transformer.py

Line 77 in 742c753

return hs

However, in detr_vae.py, the action_head uses only hs[0] as input. This means the action is predicted based only on the first decorder feature with shape [1, B, num_query, C], but ignores all the other following features:

act/detr/models/detr_vae.py

Line 131 in 742c753

hs = self.transformer(src, None, self.query_embed.weight, pos, latent_input, proprio_input, self.additional_pos_embed.weight)[0]

act/detr/models/detr_vae.py

Line 136 in 742c753

hs = self.transformer(transformer_input, None, self.query_embed.weight, self.pos.weight)[0]

So, this is strange for me and I wonder if this is a bug. Hoping for your reply! Thanks in advance~

[Facebear-ljx]

I checked the official DETR code and found the transfomer in DETR returns a tuple [hs, memory]:

https://github.com/facebookresearch/detr/blob/29901c51d7fe8712168b8d0d64351170bc0f83e0/models/transformer.py#L59

So, hs[0] in their code uses all intermediate features rather than the first decorder feature.