(model, dtype = torch.float32)
| 283 | } |
| 284 | |
| 285 | def get_dit_module(model, dtype = torch.float32): |
| 286 | dit_model = model.model |
| 287 | dit_model = dit_model.to(dtype).eval().requires_grad_(False) |
| 288 | return dit_model |
| 289 | |
| 290 | |
| 291 | ## Utility functions for AutoEncoder |