Thanks @ptrblck. Basically, I implemented my own class following the keras schema, using the transformers hugging face package
here is the internal part of the code in fit() function for one epoch
losses = []
predictions = []
 for step, batch in  enumerate(data_loader):
            input_ids = batch["input_ids"].to(self.device)
            attention_mask = batch["attention_mask"].to(self.device)
            targets = batch["targets"].to(self.device)
            outputs = self.__call__(
                input_ids=input_ids,
                attention_mask=attention_mask
            )
            loss = self.loss_func(outputs, targets)
           ......
           with amp.scale_loss(loss, self.optimizer) as scaled_loss:
                     scaled_loss.backward()
            nn.utils.clip_grad_norm_(self.parameters(), max_norm=1.0)
            self.optimizer.step()
            self.scheduler.step()
            self.optimizer.zero_grad()
return np.mean(predictions), np.mean(losses)

Alos pytorch provide an other way to quantize the model.  I have used this with bertdistill trained with mixed precision. The quantized model performs as good as the float32 version!
So my though is that the mixed precision is used just for the training phase and the model is saved all as 32 float??
model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)