@ptrblck this is how I am doing the forward pass(calculating the loss):
def cal_loss(unrolled_model,unrolled_feature_extractor,unrolled_head_g,input_valid_source,input_valid_target,target_valid_source,target_valid_target):
domain_features_source,domain_logits_s=unrolled_model(input_valid_source)
domain_features_target,domain_logits_t=unrolled_model(input_valid_target)
generalized_features_source = unrolled_feature_extractor(input_valid_source)
generalized_features_target = unrolled_feature_extractor(input_valid_target)
logits1=unrolled_head_g(generalized_features_source-domain_features_source)
logits2=unrolled_head_g(generalized_features_target-domain_features_target)
crit = nn.CrossEntropyLoss()
loss1 = crit(logits1, target_valid_source)
loss2 = crit(logits2, target_valid_target)
loss=loss1+loss2
return loss