Returning a trained model from a class?

1

It is more of a python question than pytorch. I have a model which I would like to pass to class A object, then the same model should be returned from class A object, and passed to class B object.

Class A and Class B have different data to train the model, but instead of passing the model independently to class A and then class B, I would like to pass returned model from class A to class B.

If anyone could help or suggest material that would be amazing.
Thank you

2

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))