different method of running pytorch on gpu
See the code block below (the source of the code can be found here, also you don't need to read the whole block, I will explain and highlight the important part)
def train(data_loader, model, optimizer, scheduler, total_epochs, save_interval, save_folder, sets):
    # settings
    batches_per_epoch = len(data_loader) # 
    log.info('{} epochs in total, {} batches per epoch'.format(total_epochs, batches_per_epoch))
    loss_seg = nn.CrossEntropyLoss(ignore_index=-1) 

    print(&quot;Current setting is:&quot;)
    print(sets)
    print(&quot;\n\n&quot;)     
    if not sets.no_cuda:
        loss_seg = loss_seg.cuda()
    
    # change model in training mode, enable batch normalization etc
    model.train()
    
    # record train time 
    train_time_sp = time.time()
    
    # loop to train the model
    for epoch in range(total_epochs):
        log.info('Start epoch {}'.format(epoch))
        
        scheduler.step()
        log.info('lr = {}'.format(scheduler.get_lr()))
        
        for batch_id, batch_data in enumerate(data_loader):
            # getting data batch
            batch_id_sp = epoch * batches_per_epoch
            volumes, label_masks = batch_data

            if not sets.no_cuda: 
                volumes = volumes.cuda()

            optimizer.zero_grad()
            out_masks = model(volumes)
            # resize label
            [n, _, d, h, w] = out_masks.shape
            new_label_masks = np.zeros([n, d, h, w])
            for label_id in range(n):
                label_mask = label_masks[label_id]
                [ori_c, ori_d, ori_h, ori_w] = label_mask.shape 
                label_mask = np.reshape(label_mask, [ori_d, ori_h, ori_w])
                scale = [d*1.0/ori_d, h*1.0/ori_h, w*1.0/ori_w]
                label_mask = ndimage.interpolation.zoom(label_mask, scale, order=0)
                new_label_masks[label_id] = label_mask

            new_label_masks = torch.tensor(new_label_masks).to(torch.int64)
            if not sets.no_cuda:
                new_label_masks = new_label_masks.cuda()

            # calculating loss
            loss_value_seg = loss_seg(out_masks, new_label_masks)
            loss = loss_value_seg
            loss.backward()                
            optimizer.step()

            avg_batch_time = (time.time() - train_time_sp) / (1 + batch_id_sp)
            log.info(
                    'Batch: {}-{} ({}), loss = {:.3f}, loss_seg = {:.3f}, avg_batch_time = {:.3f}'\
                    .format(epoch, batch_id, batch_id_sp, loss.item(), loss_value_seg.item(), avg_batch_time))
          
            if not sets.ci_test:
                # save model
                if batch_id == 0 and batch_id_sp != 0 and batch_id_sp % save_interval == 0:
                #if batch_id_sp != 0 and batch_id_sp % save_interval == 0:
                    model_save_path = '{}_epoch_{}_batch_{}.pth.tar'.format(save_folder, epoch, batch_id)
                    model_save_dir = os.path.dirname(model_save_path)
                    if not os.path.exists(model_save_dir):
                        os.makedirs(model_save_dir)
                    
                    log.info('Save checkpoints: epoch = {}, batch_id = {}'.format(epoch, batch_id)) 
                    torch.save({
                                'ecpoch': epoch,
                                'batch_id': batch_id,
                                'state_dict': model.state_dict(),
                                'optimizer': optimizer.state_dict()},
                                model_save_path)
                            
    print('Finished training')            
    if sets.ci_test:
        exit()

This is a customised training function, the author achieve GPU running by using code loss_seg = loss_seg.cuda() where loss_seg is a instance of torch.optim.SGD. This part confused me, because according to the official documentation in pytorch, I only need to move my model and my input data to gpu is enough. I want to know why the author who write the code below also move the optimizer to the gpu and more details about running pytorch on gpu.

In general, in order to harness the full power of GPUs, every stateful Module should be sent to a cuda device before the forward step. A stateful Module has an internal state, e.g. Parameter (weights).
This is not usually the case of Loss, which, in general, applies just a functional that has been already implemented for cuda devices.
In conclusion, if the loss is stateful, then it makes sense to send it to a cuda device, otherwise, it is not necessary.
Have also a look at this question in the PyTorch forum.
