(model, dtype = torch.float32)
| 269 | } |
| 270 | |
| 271 | def get_dit_module(model, dtype = torch.float32): |
| 272 | dit_model = model.model |
| 273 | dit_model = dit_model.to(dtype).eval().requires_grad_(False) |
| 274 | return dit_model |
| 275 | |
| 276 | |
| 277 | ## ----------------- Utility Functions AutoEncoder ------------------- |