I’m not sure, what your use case is, but would this work?
class classA(object):
    def train(self, model):
        # train model
        return model

class classB(object):
    def train(self, model):
        # train model
        return model

model = nn.Linear(10, 10)
A = classA()
B = classB()
model = A.train(model)
model = B.train(model)

# alternatively
model = B.train(A.train(model))