from cutblur.augments import cutblur
def train_loop_fn(data_loader, model, optimizer, device, scheduler=None):
    running_loss = 0.0
    running_corrects = 0
    
    model.train()
    
    for inputs,labels in data_loader:
        inputs_HQ = inputs.to(device, dtype=torch.float)
        labels = labels.to(device, dtype=torch.float)

        # or you can apply random noise, jittering, etc..
        inputs_LQ = F.interpolate(inputs_HQ, scale_factor=1/4, mode="bilinear")
        inputs = cutblur(inputs_HQ, inputs_LQ)
        
        optimizer.zero_grad()

        outputs = model(inputs)
        _, preds = torch.max(outputs, 1)
        
        loss = loss_fn(outputs, label)

        loss.backward()
        xm.optimizer_step(optimizer)

        running_loss += loss.item()
        running_corrects += torch.sum(preds == label.data)
            
    train_loss = running_loss / float(len(train_data))
    train_acc = running_corrects.double() / float(len(train_data))
    
    scheduler.step(train_loss)
    
    xm.master_print('training Loss: {:.4f} & training accuracy : {:.4f}'.format(train_loss , train_acc))