class DCTLN(nn.Cell):
    def __init__(self, input_shape, num_classes, BatchNormal):
        super().__init__()
        self.feature_extractor = FeatureExtractor(input_shape, 20, 5, 1, 'same', 2, BatchNormal)
        self.hd_classifier = HealthConditionClassifier(100, num_classes)
        self.domain_classifier = DomainClassifier(100)
        self.grl = GradReverse()
        self.apply(self._init_weights)
        
    def construct(self, x):
        features = self.feature_extractor(x)
        health_condition = self.hd_classifier(features)
        domain = self.domain_classifier(self.grl(features))
        return features, health_condition, domain

def forward_fn(X_source, y_source, X_target, y_target, penalty_param):
    X = ops.concat([X_source, X_target], axis=0)
    # predict and calculate the loss
    features, logits, domain = model(X)
    
    classification_loss = loss_fn(logits.narrow(0, 0, source_size), y_source)
    domain_label = ops.concat([ops.zeros(source_size), ops.ones(target_size)], axis=0)
    domain_loss = ops.binary_cross_entropy(domain, domain_label.reshape(-1, 1))

    mmd_value = mmd(features.narrow(0, 0, source_size), features.narrow(0, source_size, target_size))
    
    # 更新惩罚系数
    model.set_lambda(penalty_param)
    loss = classification_loss #+ domain_loss + penalty_param * mmd_value
    
    return loss, logits
    
grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
(loss, logits), grads = grad_fn(X_source, y_source, X_target, y_target, penalty_param)
# backpropagation
optimizer(grads)

Model(net, loss, opt, metrics={'acc', 'loss'})