Hello all,
I am trying to convert my model to half-precision followed by mixed-precision training so that lesser memory could be used. I am not sure whether this is the correct way to do it -

    def train_model(self, model, dataloader, epochs):
        cudnn.benchmark = True
        optimizer = torch.optim.SGD(model.parameters(),lr=0.01)
        model.half()
        model ,optimizer = amp.initialize(model, optimizer, opt_level='01')
        model.train()
        model.cuda()
        scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer, step_size=10, gamma=0.1)
        #criterion = torch.nn.CrossEntropyLoss().cuda()
        criterion = torch.nn.BCELoss().cuda()
        for i in range(0, epochs):
            scheduler.step()
            train_accuracy = 0
            net_loss = 0
            for _, (data, label) in enumerate(dataloader):
                optimizer.zero_grad()
                data = data.half().cuda()
                label = label.cuda()
                out = model(data)
                loss = criterion(out, label)
                with amp.scale_loss(loss,optimizer) as scaled_loss:
                    scaled_loss.backward()
                    optimizer.step()
                if torch.argmax(out) == label:
                    train_accuracy +=1
                net_loss += loss.item()
            print('------------------------------------------')
            print('EPOCH ', i)
            print(train_accuracy/len(dataloader))
            print(net_loss/len(dataloader))

Could someone please guide me in the correct direction
Thanks in Advance