diff --git a/studio/backend/tests/test_diffusion_dit_trainer.py b/studio/backend/tests/test_diffusion_dit_trainer.py
index 6d7736a106..d406bf99eb 100644
--- a/studio/backend/tests/test_diffusion_dit_trainer.py
+++ b/studio/backend/tests/test_diffusion_dit_trainer.py
@@ -236,6 +236,7 @@ def _patch_capability(monkeypatch, capability):
     import torch
 
     import core.training.diffusion_train_common as dtc
+
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
     monkeypatch.setattr(torch.cuda, "get_device_capability", lambda *a, **k: capability)
     monkeypatch.setattr(dtc, "has_functional_torchao", lambda: True)