Hello, I’m trying to train Neural Networks using format datatype BFloat16 in Pytorch.
I’ve started with a simple example. I’ve tried to train LeNet5 with MNIST dataset.
Firstly, I’ve extracted the datasets and dataloaders with the next code:
transforms = transforms.Compose([transforms.Resize((32, 32)),
                                 transforms.ToTensor(),
                                 transforms.ConvertImageDtype(dtype= torch.bfloat16)]) 

# download and create datasets
train_dataset = datasets.MNIST(root='mnist_data', 
                               train=True, 
                               transform=transforms,
                               download=True) 

valid_dataset = datasets.MNIST(root='mnist_data', 
                               train=False, 
                               transform=transforms) 

# define the data loaders
train_loader = DataLoader(dataset=train_dataset, 
                          batch_size=BATCH_SIZE, 
                          shuffle=True)  

valid_loader = DataLoader(dataset=valid_dataset, 
                          batch_size=BATCH_SIZE, 
                          shuffle=False) 

Apparently, this works, and the dataset has Bfloat16 datatype.
The problem I have is when I try to execute the model.
This is my train function:
def train(train_loader, model, criterion, optimizer, device):
    model.train() 
    running_loss = 0 
    
    for X, y_true in train_loader: 

        optimizer.zero_grad() 
      
        X = X.to(device) 
        y_true = y_true.to(device) 
        
        # Forward pass
        y_hat, _ = model(X) 
        loss = criterion(y_hat, y_true) 
        running_loss += loss.item() * X.size(0) 

        # Backward pass
        loss.backward() 
        optimizer.step() 
        
    epoch_loss = running_loss / len(train_loader.dataset)
    return model, optimizer, epoch_loss 

An, when I execute the model function. I get the error.
This is how I configure the model execution:
torch.manual_seed(RANDOM_SEED) 
model = LeNet5(N_CLASSES).to(DEVICE) 
optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE) 
criterion = nn.CrossEntropyLoss() 

model, optimizer, _ = training_loop(model, criterion, optimizer, train_loader, valid_loader, N_EPOCHS, DEVICE)

And this is the error I get:
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-45-d0a5e863d7fc> in <module>
----> 1 model, optimizer, _ = training_loop(model, criterion, optimizer, train_loader, valid_loader, N_EPOCHS, DEVICE)

<ipython-input-37-247a8affb032> in training_loop(model, criterion, optimizer, train_loader, valid_loader, epochs, device, print_every)
     13 
     14         # training
---> 15         model, optimizer, train_loss = train(train_loader, model, criterion, optimizer, device) #Llamada a la función de entrenamiento
     16         train_losses.append(train_loss) #Añadimos al array de pérdidas de entrenamiento, la périda del entrenamiento de la época
     17 

<ipython-input-35-c9cf685f7d11> in train(train_loader, model, criterion, optimizer, device)
     18 
     19         # Forward pass
---> 20         y_hat, _ = model(X) #Aplicamos el modelo en X y obtenemos la predicción
     21         loss = criterion(y_hat, y_true) #Calculamos la pérdida comparando la predicción con el valor correcto
     22         running_loss += loss.item() * X.size(0) #Aumentamos el valor de la pérdida total sumando la pérdida actual con el tamaño de la imagen actual

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

<ipython-input-38-7278d3f7d36a> in forward(self, x)
     23 
     24     def forward(self, x): #Implementación de Lenet5
---> 25         x = self.feature_extractor(x) #Extrae las características de X y las guarda en X
     26         x = torch.flatten(x, 1) #Aplana un rango de atenuaciones en un tensor. Comienzo de dimensión: X, final de dimensión:1
     27         logits = self.classifier(x) #Clasificamos los valores de X

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/container.py in forward(self, input)
    117     def forward(self, input):
    118         for module in self:
--> 119             input = module(input)
    120         return input
    121 

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py in forward(self, input)
    397 
    398     def forward(self, input: Tensor) -> Tensor:
--> 399         return self._conv_forward(input, self.weight, self.bias)
    400 
    401 class Conv3d(_ConvNd):

~/anaconda3/lib/python3.8/site-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight, bias)
    393                             weight, bias, self.stride,
    394                             _pair(0), self.dilation, self.groups)
--> 395         return F.conv2d(input, weight, bias, self.stride,
    396                         self.padding, self.dilation, self.groups)
    397 

RuntimeError: Expected object of scalar type BFloat16 but got scalar type Float for argument #2 'weight' in call to _thnn_conv2d_forward

How can I solve this?
Thank you