#%%
import torch
from torchvision.models import resnet18
from bigdl.nano.pytorch import InferenceOptimizer

#%%
USE_JIT = True
# USE_JIT = False
# PRECISION = "fp32"
PRECISION = "bf16"

model_ft = resnet18(pretrained=True)

if PRECISION == "fp32":
    jit_model = InferenceOptimizer.trace(
        model_ft,
        accelerator="jit",
        use_ipex=True,
        input_sample=torch.rand(1, 3, 224, 224),
        thread_num=1,
    )
elif PRECISION == "bf16":
    jit_model = InferenceOptimizer.quantize(
        model_ft,
        precision="bf16",
        accelerator="jit",
        use_ipex=True,
        input_sample=torch.rand(1, 3, 224, 224),
        thread_num=1,
    )
else:
    print("Error: The PRECISION must be 'fp32' or 'bf16'")

x = torch.rand(3, 3, 224, 224)
#%%
with torch.no_grad():
    with torch.jit.optimized_execution(USE_JIT):
        %time y_hat = jit_model(x)

#%%
with torch.no_grad():
    with torch.jit.optimized_execution(USE_JIT):
        %time y_hat = jit_model(x)
#%%
with torch.no_grad():
    with torch.jit.optimized_execution(USE_JIT):
        %time y_hat = jit_model(x)
#%%
with torch.no_grad():
    with torch.jit.optimized_execution(USE_JIT):
        %time y_hat = jit_model(x)
#%%
with torch.no_grad():
    with torch.jit.optimized_execution(USE_JIT):
        %time y_hat = jit_model(x)
#%%
with torch.no_grad():
    with torch.jit.optimized_execution(USE_JIT):
        %time y_hat = jit_model(x)
#%%