Skip to content

Commit a442a21

Browse files
authored
Apply suggestions from code review
1 parent 172cb97 commit a442a21

File tree

2 files changed

+2
-3
lines changed

2 files changed

+2
-3
lines changed

tests/pipelines/controlnet_flux/test_controlnet_flux.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,8 +230,7 @@ def test_canny(self):
230230
text_encoder_2=None,
231231
controlnet=controlnet,
232232
torch_dtype=torch.bfloat16,
233-
)
234-
pipe.enable_model_cpu_offload(device=torch_device)
233+
).to(torch_device)
235234
pipe.set_progress_bar_config(disable=None)
236235

237236
generator = torch.Generator(device="cpu").manual_seed(0)

tests/pipelines/deepfloyd_if/test_if_superresolution.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,8 +127,8 @@ def test_if_superresolution(self):
127127
pipe.enable_model_cpu_offload(device=torch_device)
128128

129129
# Super resolution test
130-
backend_reset_max_memory_allocated(torch_device)
131130
backend_empty_cache(torch_device)
131+
backend_reset_max_memory_allocated(torch_device)
132132
backend_reset_peak_memory_stats(torch_device)
133133

134134
image = floats_tensor((1, 3, 64, 64), rng=random.Random(0)).to(torch_device)

0 commit comments

Comments
 (0)