//  [python]
def compute_loss(self, X, Y, I0, idx):
        ep = 1e-7 # to be used if a normalized loss is considered (check the variable "res")
        
        p0 = (gamma2-gamma1)*self.sigmoid_in(self.gamma[0,])+gamma1
        p1 = (delta2-delta1)*self.sigmoid_in(self.delta[0,])+delta1
        
        E0l = 3*I0*self.sigmoid_in(self.E0[idx,:,:])
        S0 = 1-I0-E0l
        
        # compute I(t) from R(t)+I(t)
        bIA = torch.cat( ( (1-(self.Ntf-1)*self.dt/2*p0)*I0 - Y[:,0:1,None] , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1)
        
        IAm = torch.linalg.solve_triangular(self.expli + (self.Ntf-1)*self.dt*p0*self.trap, bIA + torch.matmul( self.expli , Y[:,1:,None] ), upper=False)
        If = torch.cat( ( I0 , IAm ), 1)

        beta  = self.forward_pass(X, self.Nt)
        
        # Build RHS of S equation
        bS = torch.cat( ( S0*( 1 - (self.Ntf-1)*self.dt/2*beta[:,0:1,:]*If[:,0:1,:] ) , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1) # size: batch_county x Nt-1 x 1 
        # Build matrix of S equation
        betaIm = self.matrix_diag ( (beta[:,1:,:]*If[:,1:,:])[:,:,0] ) # size: batch_county x Nt-1 x Nt-1
        betaIm1 = self.matrix_diag( (beta[:,1:-1,:]*If[:,1:-1,:])[:,:,0] ) # size: batch_county x Nt-2 x Nt-2
        betaIm1f = torch.cat( (torch.zeros(self.batch_county,1,self.Nt-2).type(self.dtype),betaIm1) , 1 ) # size: batch_county x Nt-1 x Nt-2
        betaIm1ff = torch.cat( (betaIm1f,torch.zeros(self.batch_county,self.Nt-1,1).type(self.dtype)) , 2 ) # size: batch_county x Nt-1 x Nt-1
        # S prediction
        Sm = torch.linalg.solve_triangular(self.expli +(self.Ntf-1)*self.dt/2*(betaIm+betaIm1ff), bS, upper=False) # size: batch_county x Nt-1 x 1
        
        # Build RHS of E equation
        bE = torch.cat( ( E0l + (self.Ntf-1)*self.dt/2*( S0*beta[:,0:1,:]*If[:,0:1,:] - p1*E0l ) , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1) # size: batch_county x Nt-1 x 1 
        betaSIm = beta[:,1:,:]*If[:,1:,:]*Sm # size: batch_county x Nt-1 x 1 
        # E prediction
        Em = torch.linalg.solve_triangular(self.expli + (self.Ntf-1)*self.dt*p1*self.trap, bE + (self.Ntf-1)*self.dt*torch.matmul(self.trap,betaSIm), upper=False) # size: batch_county x Nt-1 x 1 
        
        # Build RHS of I equation
        bI = torch.cat( ( If[:,0:1,:] + (self.Ntf-1)*self.dt/2* (p1*E0l-p0*If[:,0:1,:]) , torch.zeros(self.batch_county,self.Nt-2,1).type(self.dtype) ), 1)
        if self.is_pred == 1:
            # I prediction
            pred = torch.linalg.solve_triangular(self.expli + (self.Ntf-1)*self.dt*p0*self.trap, bI + (self.Ntf-1)*self.dt*p1*torch.matmul(self.trap,Em), upper=False) # size: batch_county x Nt-1 x 1 
            res = (If[:,1:,:] - pred) #/ (torch.mean(If[:,1:,0],dim=1)[:,None,None] + ep) # size: batch_county x Nt-1 x 1 
        else:
            # I equation residual
            res = ( torch.matmul( (self.expli + (self.Ntf-1)*self.dt*p0*self.trap) , If[:,1:,:] ) - ( bI + (self.Ntf-1)*self.dt*p1*torch.matmul(self.trap,Em) ) ) #/ (torch.mean(If[:,1:,0]-If[:,:-1,0],dim=1)[:,None,None] + ep) # size: batch_county x Nt-1 x 1    
               
        loss = torch.mean(res[:,torch.randperm(self.Nt-1)[0:self.batch_size],:]**2)    
        
        return loss