fix lint

chainyo · May 3, 2023 · bb269c7 · bb269c7
1 parent 6cfbbfa
commit bb269c7
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/picaisso/api/diffusion_model.py b/picaisso/api/diffusion_model.py
@@ -26,11 +26,7 @@ def __init__(
         max_wait: int,
     ):
         self.model = model_name
-        self.dtype = {
-            "fp32": torch.float32,
-            "fp16": torch.float16,
-            "bf16": torch.bfloat16
-        }[dtype]
+        self.dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[dtype]
         self.n_steps = n_steps
         self.max_batch_size = max_batch_size
         self.max_wait = max_wait