Segmentation for full image

[aqibsaeed]

Hi,

Thank you for your effort and time in implementing this. I have a quick question, I want to get segmentation for full image not just for the middle token, would it be correct to change self.tokens to self.p here:

self-attention-cv/self_attention_cv/Transformer3Dsegmentation/tranf3Dseg.py

Line 66 in 5246e55

self.mlp_seg_head = nn.Linear(dim, self.tokens * self.num_classes)

and change this:

self-attention-cv/self_attention_cv/Transformer3Dsegmentation/tranf3Dseg.py

Line 94 in 5246e55

y = self.mlp_seg_head(y[:, self.mid_token, :])

to

y = self.mlp_seg_head(y)