I am training distilBert model for text classification. I am using the amp package to train the mixed precision version. here is the code to train the model
     # define the model
      model = BertMulticlassifier(....)

    # Define optimizer and scheduler
    optimizer = Config.OPTIMIZER(model.parameters(), lr=Config.LR, bias_correction=False)
    # model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
    scheduler = Config.SCHEDULER(optimizer=optimizer,
                                 num_warmup_steps=400,
                                 num_training_steps=len(train_data_loader) * Config.EPOCHS
                                 )

    model.compile(optimizer=optimizer, loss=loss_func, metrics=metric, scheduler=scheduler)
    model.fit(train_loader=train_data_loader, val_loader=val_data_loader, epochs=Config.EPOCHS)
   # save the model
    torch.save(self.state_dict(), model_name)

Load the model and run it on CPU:
model = BertMulticlassifier(....)
model.load_state_dict(torch.load('model.bin', map_location='cpu'))
model.eval()
model.predict()

By comparing the inference time between the distilbert and the mixed precision distilbert version, it was more or less the same, which does not make sense. I expect that the mixed precision version take less time.
So, I am wondering that the amp just used the mixed precision to make the train and then it saved the model with float32 or the problem is coming from the load function ? What is the right way to  load the model in the mixed precision version. Thanks to take a look @ptrblck