Defining python class method using arguments from __init__
What I want to achieve
I want to define a class with two modes A and B, so that the forward method of the class changes accordingly.
class MyClass():
    def __init__(self, constant):
        self.constant=constant
            
    def forward(self, x1,x2,function):
        if function=='A':
            return x1+self.constant
        
        elif function=='B':
            return x1*x2+self.constant
        
        else:
            print('please provide the correct function')
model1 = MyClass(2)
model1.forward(2, None, 'A')
output&gt;&gt;&gt;4
model2 = MyClass(2)
model2.forward(2, 2, 'B')
output&gt;&gt;&gt;6

It works, but it is not optimal, since every time when calling the forward method, it will check which function to use. However, the forward function is already set and will never be changed once the class is define, therefore, checking which function to use inside forward is super redundant in my case. (For those who notice this, I am writing my neural network model using PyTorch, two models share 90% of the network architecture, the only 10% differences is the way they do feedforward).
My desired version
I want to set the forward method when the class is defined, so that I can achieve this
model1 = MyClass(2, 'A')
model1.forward(2)
output&gt;&gt;&gt;4

model2 = MyClass(2, 'B')
model2.forward(2, 2)
output&gt;&gt;&gt;6

So I rewrote my class to be:
class MyClass():
    def __init__(self, constant, function):
        self.constant=constant # There would be a lot of shared parameters for the two methods
        self.function=function # This controls the feedforward method of this class
        
    if self.function=='A':
        def forward(self, x1):
            return x1+self.constant
        
    elif self.function=='B':
        def forward(self, x1, x2):
            return x1*x2+self.constant
        
    else:
        print('please provide the correct function')

However, it gives me the following error.

NameError: name 'self' is not defined

How do I write the class for that it defines different forward method based on the args from __init__?

You have been trying to redefine the class with your code, such that each new object would change the definition of forward for all objects, before and after.
Fortunately, you didn't figure out how to do that.
Instead, make the chosen function an attribute of the object.  Code the two functions you want, and then assign the desired variant as you create each instance.
class MyClass():
    def __init__(self, constant, function):
        self.constant=constant
        if function == 'A':
            self.forward = self.forwardA
        elif function=='B':
            self.forward = self.forwardB
        else:
            print('please provide the correct function')
            
    def forwardA(self, x1):
        return x1+self.constant
        
    def forwardB(self, x1, x2):
        return x1*x2+self.constant

# Main
model1 = MyClass(2, 'A')
print(model1.forward(2))

model2 = MyClass(2, 'B')
print(model2.forward(2, 2))

Output:
4
6

