Solved by ptrblck in post #2 
              
              
                You could pass the target tensor into the forward method and use it to skip the actual calculation of head in the last for loop, if this output is not needed (i.e. if no target is available for it). 
During validation, you would have to make sure that the target tensor is ignored and I assume you wa…