# save the model
unet_traced.save("unet_traced.pt")
```  
Replace the `unet` attribute of the pipeline with the traced model:  
```python
from diffusers import StableDiffusionPipeline
import torch
from dataclasses import dataclass


@dataclass
class UNet2DConditionOutput:
sample: torch.Tensor


pipe = StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
use_safetensors=True,
).to("cuda")