(self, closure=None, print_flag=False)
return loss
Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
Performs a single optimization step.
Performs a single optimization step .
def step(self, closure=None, print_flag=False):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            weight_decay = group['weight_decay']
            momentum = group['momentum']
            dampening = group['dampening']
            nesterov = group['nesterov']


            for p in group['params']:
                if p.grad is None:
                    continue
                d_p = p.grad.data


                if weight_decay != 0:
                    d_p.add_(weight_decay, p.data)

                sz = p.data.size()
                if d_p.dim() == 4 and sz[1] != 1: # we do not consider dw conv
                    assert(weight_decay == 0)
                    sz = p.data.size()
                    w  = p.data.view(sz[0], -1)
                    wstd = w.std(dim=1).view(sz[0], 1, 1, 1)
                    wmean = w.mean(dim=1).view(sz[0], 1, 1, 1)

                    if args.local_rank == 0 and print_flag:
                        wm = wstd.view(-1).mean().item()
                        wmm = wmean.view(-1).mean().item()
                        print('lam = %.6f' % args.lam, 'mineps = %.6f' % args.mineps, 
                                '1 - eps/std = %.10f' % (1 - args.mineps / wm), 
                                'std = %.10f' % wm, 'mean = %.10f' % wmm,  'sz = ', sz)
                    
                    d_p.add_(args.lam, (1 - args.mineps / wstd) * (p.data - wmean) + wmean)


                if momentum != 0:
                    param_state = self.state[p]
                    if 'momentum_buffer' not in param_state:
                        buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
                        buf.mul_(momentum).add_(d_p)
                    else:
                        buf = param_state['momentum_buffer']
                        buf.mul_(momentum).add_(1 - dampening, d_p)
                    if nesterov:
                        d_p = d_p.add(momentum, buf)
                    else:
                        d_p = buf

                p.data.add_(-group['lr'], d_p)

        return loss
def step ( self , closure = None , print_flag = False ) : loss = None if closure is not None : loss = closure ( ) for group in self . param_groups : weight_decay = group [ 'weight_decay' ] momentum = group [ 'momentum' ] dampening = group [ 'dampening' ] nesterov = group [ 'nesterov' ] for p in group [ 'params' ] : if p . grad is None : continue d_p = p . grad . data if weight_decay != 0 : d_p . add_ ( weight_decay , p . data ) sz = p . data . size ( ) if d_p . dim ( ) == 4 and sz [ 1 ] != 1 : # we do not consider dw conv assert ( weight_decay == 0 ) sz = p . data . size ( ) w = p . data . view ( sz [ 0 ] , - 1 ) wstd = w . std ( dim = 1 ) . view ( sz [ 0 ] , 1 , 1 , 1 ) wmean = w . mean ( dim = 1 ) . view ( sz [ 0 ] , 1 , 1 , 1 ) if args . local_rank == 0 and print_flag : wm = wstd . view ( - 1 ) . mean ( ) . item ( ) wmm = wmean . view ( - 1 ) . mean ( ) . item ( ) print ( 'lam = %.6f' % args . lam , 'mineps = %.6f' % args . mineps , '1 - eps/std = %.10f' % ( 1 - args . mineps / wm ) , 'std = %.10f' % wm , 'mean = %.10f' % wmm , 'sz = ' , sz ) d_p . add_ ( args . lam , ( 1 - args . mineps / wstd ) * ( p . data - wmean ) + wmean ) if momentum != 0 : param_state = self . state [ p ] if 'momentum_buffer' not in param_state : buf = param_state [ 'momentum_buffer' ] = torch . zeros_like ( p . data ) buf . mul_ ( momentum ) . add_ ( d_p ) else : buf = param_state [ 'momentum_buffer' ] buf . mul_ ( momentum ) . add_ ( 1 - dampening , d_p ) if nesterov : d_p = d_p . add ( momentum , buf ) else : d_p = buf p . data . add_ ( - group [ 'lr' ] , d_p ) return loss