great!
another question, hope it’s alright  
i notice my gpu utilization is declining rapidly thanks to this embedding procedure, as it causes some bottleneck along the way. can you think of any method of reducing it?
my actual input data is of size (batch=4,channel=1,width=75,height=95,depth=65) -
this means my self.embed is of size (category_num=155,channel=1,width=75,height=95,depth=65)
and so is self.mask.
the init func is now as so:
        if args.subject_embedding:
            self.do_embed = True
            self.embed = torch.randn((self.n_subjects, 1, self.dim[0], self.dim[1], self.dim[2]), requires_grad=True)
            self.mask = torch.ones(self.embed.shape,requires_grad=False)
            if args.embedding_mask is not None:
                self.mask = args.embedding_mask
                self.mask.requires_grad = False
                if self.mask.shape[0] == 1:
                    self.mask = self.mask.unsqueeze(0).repeat(self.n_subjects, 1, 1, 1, 1)
                else:
                    self.mask = self.mask.unsqueeze(0).unsqueeze(0).repeat(self.n_subjects, 1, 1, 1, 1)
            if args.cuda:
                self.mask = self.mask.cuda()
                self.embed = self.embed.cuda()

the forward pass is implemented as so:
    def forward(self, x, subj):

        #  Level 1 context pathway
        #print('input shape ',x.shape)
        context_0 = x
        if self.do_embed:
            embed = self.embed * self.mask
            E = embed[subj.to(dtype=torch.long)]
            x = torch.cat([x,E],dim=1)

still, since the only nonzero elements of self.mask are 40 per category, i thought this would reduce the bottleneck. but now i see it doesn’t and training will be slow.
notice that i push both the entire huge embedding to cuda and the huge ones tensor to cuda, where effectively i’m not sure that is necessary.
i was wondering if perhaps keeping them in cpu and only pushing to cuda the result of their multiplication (the tensor E ) in the forward pass, would reduce the bottleneck. i implemented it and couldn’t prove myself it’s more efficient yet.
P.S
do not be confused by the initiation of self.mask as a ones tensor, that is for the case the masking feature is turned off and then it operates as an identity transformation.