My hypothesis of writing customized loss function to deal with customized model was wrong Then I tried something different.
I have figured out one solution using this reference 5 (but again there are some issues in the model I am not able to understand why this is happening)
My codes are as follows:
class CustomizedForwardBackwardFunc(torch.autograd.Function):
    '''
        Autograd function which masks it's weight by 'mask'
    '''
    
    @staticmethod
    def forward(ctx, input, production_weight, decay_weight, 
                cross_talk_weight, bias, prod_mask, decay_mask, cross_mask):
        
        if prod_mask.size() == production_weight.size():
            production_weight = torch.mul(production_weight, prod_mask)
            decay_weight = torch.mul(decay_weight, decay_mask)
            cross_talk_weight = torch.mul(cross_talk_weight, cross_mask)
        else:            
            raise Exception('Please check your production mask matrix size')
            
        xy_term = input[0][0] * input[0][2]        
        input_new = torch.cat((input,xy_term.view([-1,1])),1)
        input_new = input_new.to(device,non_blocking=True)   
        
        production_dt = production_weight.mm(torch.transpose(input_new,0,1))
        decay_dt = decay_weight.mm(torch.transpose(input_new,0,1))
        cross_talk_dt = cross_talk_weight.mm(torch.transpose(input_new,0,1))
        hidden_state = production_dt - decay_dt - cross_talk_dt
        out = hidden_state.view([-1,3]) + input
        
        if bias is not None:
            out += bias.unsqueeze(0).expand_as(out)
            
        ctx.save_for_backward(input_new,production_weight,decay_weight,
                              cross_talk_weight,bias,prod_mask,decay_mask,
                              cross_mask)
        
        return out
    
    
    @staticmethod
    def backward(ctx, grad_out):
                
        input_new,production_weight,decay_weight,cross_talk_weight,bias,prod_mask,decay_mask,cross_mask = ctx.saved_tensors
        grad_input = grad_prod_wt = grad_decay_wt = grad_cross_wt = None
        grad_bias = grad_prod_mask = grad_decay_mask = grad_cross_mask = None        
        
        grad_out_new = torch.cat((grad_out,(grad_out[0][0] + grad_out[0][2]).view([-1,1])),1)
        
        if ctx.needs_input_grad[0]:            
            grad_input = (production_weight - decay_weight - cross_talk_weight).mm(grad_out_new)
        
        if ctx.needs_input_grad[1]:
            grad_prod_wt = grad_out.t().clone().mm(input_new)            
            grad_prod_wt = torch.mul(grad_prod_wt, prod_mask)

            grad_decay_wt = grad_out.t().clone().mm(input_new)
            grad_decay_wt = torch.mul(grad_decay_wt, decay_mask)

            grad_cross_wt = grad_out.t().clone().mm(input_new)
            grad_cross_wt = torch.mul(grad_cross_wt, cross_mask)
                        
#         if ctx.needs_input_grad[2]: #don't train for bias
#             grad_bias = (grad_prod_out + grad_decay_out + grad_cross_out).sum(0).squeeze(0)
            
        return grad_input, grad_prod_wt, grad_decay_wt, grad_cross_wt, grad_bias, grad_prod_mask, grad_decay_mask, grad_cross_mask
        
        
class NeuralNet(nn.Module):
    
    def __init__(self,param, prod_mask,decay_mask,cross_mask, bias=True):
        super(NeuralNet, self).__init__() 
        
        if isinstance(prod_mask, torch.Tensor):
            self.prod_mask = prod_mask.type(torch.float)
        else:
            self.prod_mask = torch.tensor(self.prod_mask,dtype=torch.float)
            
        self.prod_mask = nn.Parameter(self.prod_mask.to(device),requires_grad=False)
        
        if isinstance(decay_mask, torch.Tensor):
            self.decay_mask = decay_mask.type(torch.float)
        else:
            self.decay_mask = torch.tensor(self.decay_mask,dtype=torch.float)
            
        self.decay_mask = nn.Parameter(self.decay_mask.to(device),requires_grad=False)
        
        if isinstance(cross_mask, torch.Tensor):
            self.cross_mask = cross_mask.type(torch.float)
        else:
            self.cross_mask = torch.tensor(self.cross_mask,dtype=torch.float)
            
        self.cross_mask = nn.Parameter(self.cross_mask.to(device),requires_grad=False)
        
        if bias:
            self.bias = Variable(torch.zeros(self.prod_mask.shape[0])).type(dtype)
            self.bias = nn.Parameter(self.bias.to(device), requires_grad=False)
        else:
            self.register_parameter('bias',None)
        
        # param_order = [gamma,alpha_xy,beta_y,alpha0,alpha_y,alpha1,alpha2,alpha3]
        # Production Matrix
        self.w_production = torch.tensor([[param[0], 0.0, 0.0, 0.0],
                                          [param[2], 0.0, 0.0, 0.0],
                                          [0.0, param[3], 0.0, 0.0]])
        self.w_production = Variable(self.w_production).type(dtype)
        self.w_production = nn.Parameter(self.w_production.to(device),
                                         requires_grad=True)
        
        # Degradation Matrix
        self.w_decay = torch.tensor([[param[5], 0.0, 0.0, 0.0],
                                     [0.0, param[6], 0.0, 0.0],
                                     [0.0, 0.0, param[4], 0.0]])
        self.w_decay = Variable(self.w_decay).type(dtype)
        self.w_decay = nn.Parameter(self.w_decay.to(device),
                                    requires_grad = True)        
        
        # Cross-Talk Matrix
        self.w_cross_talk = torch.tensor([[0.0, 0.0, 0.0, param[1]],
                                          [0.0, 0.0, 0.0, 0.0],
                                          [0.0, 0.0, 0.0, param[7]]])
        self.w_cross_talk = Variable(self.w_cross_talk).type(dtype)
        self.w_cross_talk = nn.Parameter(self.w_cross_talk.to(device),
                                         requires_grad = True) 
               
         
        self.w_production.data = torch.mul(self.w_production.data, self.prod_mask)
        self.w_decay.data = torch.mul(self.w_decay.data, self.decay_mask)
        self.w_cross_talk.data = torch.mul(self.w_cross_talk.data, self.cross_mask)
        
        
        
    def forward(self,input): 
        
        return CustomizedForwardBackwardFunc.apply(input,self.w_production,
                                               self.w_decay,self.w_cross_talk,
                                               self.bias,self.prod_mask,
                                               self.decay_mask,
                                               self.cross_mask)
    
    

# params initialized  from lmfit
# init_params = [gamma,alpha_xy,beta_y,alpha0,alpha_y,alpha1,alpha2,alpha3]
# init_params = [2.2, 3.1, 1.1, 1.0, 0.79, 0.07, 0.5, 0.008]
init_params = [2.0,3.7,1.5,1.1,0.9,0.1,0.9,0.01]
prod_mask = torch.tensor([[1.0, 0.0, 0.0, 0.0],
                          [1.0, 0.0, 0.0, 0.0],
                          [0.0, 1.0, 0.0, 0.0]])

decay_mask = torch.tensor([[1.0, 0.0, 0.0, 0.0],
                           [0.0, 1.0, 0.0, 0.0],
                           [0.0, 0.0, 1.0, 0.0]])

cross_mask = torch.tensor([[0.0, 0.0, 0.0, 1.0],
                           [0.0, 0.0, 0.0, 0.0],
                           [0.0, 0.0, 0.0, 1.0]])

# define model in GPU
model = NeuralNet(init_params,prod_mask,decay_mask,cross_mask)
model = model.to(device)

loss_function = nn.MSELoss()

optimizer = torch.optim.SGD(model.parameters(), lr=1e-3,momentum = 0.9)

start = datetime.now()
learning_rate = 1e-6
loss_1 = []

for i in tqdm(range(epochs)):
    total_loss = 0
    
    for j in range(x_train.size(0)):
        optimizer.zero_grad() 
        
#         for param in model.parameters():
#             if param.requires_grad == True:
#                 param.retain_grad()
#                 param.grad.data.zero_() #commented because of error of  
                                          #getting "None" from param.grad
        
        x1 = x_train[j:(j+1)]
        y1 = y_train[j:(j+1)]         
        x1 = x1.to(device,non_blocking=True)
        y1 = y1.to(device,non_blocking=True)                
        pred = model(x1)
               
        loss = loss_function(pred,y1)
        total_loss += loss               
        loss.backward()            
#         optimizer.step() # commented because we're using customized 
                           # gradient upgradation
        with torch.no_grad():
            for name, param in model.named_parameters():                
                if param.requires_grad == True:
                    param -= learning_rate * param.grad
                                        
                    '''
                        Print statement to cross-cross check the masked
                        gradradients
                    '''
#                     print(param.grad)
                    # check masked param.grad
#                     if np.array(param.grad.cpu()).size == np.array(prod_mask).size:
#                         print('--- epoch={}, loss={} ---'.format(t,loss.item()))
#                         print('↓↓↓masked weight↓↓↓')
#                         print(param)
#                         print('↓↓↓masked grad of weight↓↓↓')
#                         print(param.grad)
             
    
    # Early Stopping
    loss_1.append(total_loss.item())  
#     print(total_loss.item())
    if len(loss_1) > 1 and np.abs(loss_1[-2] - total_loss.item()) <= 0.01:
        print('Early Stopping due to dIfference in error is less than threshold (0.10)')
        print('Stopping training at %d epochs with total loss = %d' %(i,total_loss))
        break


    if i % 50 == 0:
        print("Epoch: {} loss {}".format(i, total_loss))
        
end = datetime.now()
time_taken = end - start
print('Execution Time: ',time_taken)

First I tried with learning_rate = 1e-3 loss was continuously increasing. Then I tried 1e-4 still the same model response. Then I tried 1e-5 then model loss decreases to 14 epochs but then started increasing.  So I further lowered down the learning rate to 1e-6 then the model was performing well. I also used early stopping criteria of difference in error as 0.1
Here is the model prediction results:

result1_original_param877×252 33.6 KB

My issues are as follows:

With learning rate 1e-6 and early stopping criteria (i.e. the difference in epoch error less than 0.1) then error lowered up to 1880 and then model training stopped. (Ideally, it must go down to zero). Why??
I lowered the early stopping condition to 0.01 with the same learning rate 1e-6 then the model performed well till 150 epochs (error at 150 epoch = 1880) and then error started increasing again.

I though lowering the error criteria will help the model but opposite is happening. My doubt is :

Is there anything wrong in the model that error is not going below 1880.
I am taking just one batch of all training samples. How does batch_size affect model training perofrmance.