Hi, I am trying to use the model.half() to get rid of out of memory error. So, first I want to check it for a simple model and then try it for a complex model.
In the following example, I am using model.half()for a simple CNN and MNIST dataset. I think I did all thing correctly, but I got Nan for loss.
I am wondering if you kindly let me know where I am doing a mistake.
mnist = torchvision.datasets.MNIST('./data',train= True,download=True,transform =transforms.ToTensor())
data_loader = DataLoader(mnist,batch_size=20,num_workers=2,shuffle=True)

class Model(nn.Module):
    # Our model

    def __init__(self):
        super(Model, self).__init__()
        
        self.fc1 = nn.Conv2d(1,10,3)
        self.bn1 = nn.BatchNorm2d(10)
        self.fc2= nn.Conv2d(10,20,3)
        self.bn2 = nn.BatchNorm2d(20)
        self.fc3= nn.Linear(11520,10)
        
    def forward(self,x):
        
        x = F.relu(self.fc1(x))
        
        x = x.type(torch.float32)   
        x = self.bn1(x)
                
        x = x.type(torch.float16)
        
        x = F.relu(self.fc2(x))
        
        x = x.type(torch.float32)
        
        x = self.bn2(x)
        
        
        x = x.type(torch.float16)
        
        x = x.view(x.size(0),-1)
                
        x = self.fc3(x)
        return(x)

device = torch.device('cuda:6' if torch.cuda.is_available() else 'cpu')

model = Model().half().to(device)

for layer in model.modules():
    
    if isinstance(layer,nn.BatchNorm2d):
        
        layer.float()

optimizer = optim.Adam(model.parameters(),lr=0.1)

lr_sch = lr_scheduler.StepLR(optimizer,step_size=2,gamma=0.1)

criterion = nn.CrossEntropyLoss()

def train(epoch):
    model.train()
    t_loss = 0
    for X,y in data_loader:   
        X= X.half().to(device)
        y = y.long().to(device)
        pred = model(X)
        loss = criterion(pred,y)
        t_loss+= loss.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        return(t_loss/len(data_loader.dataset))  
num_epochs = 20
train_loss = []
for epoch in range(num_epochs):
    t_loss = train(epoch)
    train_loss.append(t_loss)
    print(f'Epoch: {epoch}: Loss: {t_loss:0.4}')

Epoch: 0: Loss: 4.346e-05
Epoch: 1: Loss: nan
Epoch: 2: Loss: nan
Epoch: 3: Loss: nan
Epoch: 4: Loss: nan
Epoch: 5: Loss: nan
Epoch: 6: Loss: nan
Epoch: 7: Loss: nan
Epoch: 8: Loss: nan
Epoch: 9: Loss: nan
Epoch: 10: Loss: nan
Epoch: 11: Loss: nan
Epoch: 12: Loss: nan
Epoch: 13: Loss: nan
Epoch: 14: Loss: nan
Epoch: 15: Loss: nan
Epoch: 16: Loss: nan
Epoch: 17: Loss: nan
Epoch: 18: Loss: nan
Epoch: 19: Loss: nan