casque

/

vaeFtMse840000Ema_v100

like 4

Hey. I want to convert the model to TensorRT (https://github.com/pytorch/TensorRT) for faster inference. The code from the conversion in the docs is:

import torch
import torch_tensorrt

model = MyModel().eval().cuda()  # define your model here
x = torch.randn((1, 3, 224, 224)).cuda()  # define what the inputs to the model will look like

optimized_model = torch.compile(model, backend="tensorrt")
optimized_model(x)  # compiled on first run

optimized_model(x)  # this will be fast!

In my setup, when loading the model:

import torch
import torch_tensorrt

vaeloader = VAELoader()
vae = vaeloader.load_vae(vae_name="vaeFtMse840000Ema_v100.pt")
vae_model = vae[0]  # Get the actual VAE model from the tuple

# Ensure model is in eval mode and on CUDA
vae_model = vae_model.eval().cuda()
x = torch.randn((1, 3, 224, 224)).cuda()  # define what the inputs to the model will look like

optimized_model = torch.compile(model, backend="tensorrt")
optimized_model(x)  # compiled on first run

I am getting the following error:

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-32-16b5aa451420> in <cell line: 12>()
     10 
     11 # Ensure model is in eval mode and on CUDA
---> 12 vae_model = vae_model.eval().cuda()

AttributeError: 'VAE' object has no attribute 'eval'

Do you know what I have to do to load the VAE model in the way that the PyTorch TensorRT library is requiring it?