Set default dtype to half

2023-02-23 21:31:39 +00:00
parent de0fabbc5c
commit 1ce1333573
3 changed files with 23 additions and 3 deletions
--- a/cacheflow/worker/controller.py
+++ b/cacheflow/worker/controller.py
@@ -14,6 +14,7 @@ class Controller:
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
+        dtype: str = 'half',
    ) -> None:
        self.node_id = node_id
        self.num_workers = num_workers
@@ -35,6 +36,7 @@ class Controller:
                block_size=block_size,
                num_gpu_blocks=num_gpu_blocks,
                num_cpu_blocks=num_cpu_blocks,
+                dtype=dtype,
            )
            self.workers.append(worker)