manbeast3b commited on
Commit
7ce8553
·
verified ·
1 Parent(s): 75c3798

Update src/pipeline.py

Browse files
Files changed (1) hide show
  1. src/pipeline.py +3 -2
src/pipeline.py CHANGED
@@ -30,8 +30,9 @@ def load_pipeline() -> Pipeline:
30
  text_encoder_2=text_encoder_2,
31
  torch_dtype=dtype,
32
  )#.to("cuda")
33
- pipeline.transformer.to(memory_format=torch.channels_last)
34
- # pipeline.text_encoder.to(memory_format=torch.channels_last)
 
35
 
36
  # Optimize after moving to GPU
37
  pipeline.vae = torch.compile(pipeline.vae) # compile after moving to device
 
30
  text_encoder_2=text_encoder_2,
31
  torch_dtype=dtype,
32
  )#.to("cuda")
33
+ # pipeline.transformer.to(memory_format=torch.channels_last)
34
+ pipeline.text_encoder.to(memory_format=torch.channels_last)
35
+
36
 
37
  # Optimize after moving to GPU
38
  pipeline.vae = torch.compile(pipeline.vae) # compile after moving to device