Add utils

5e644b91 · Woosuk Kwon · d904350a · 5e644b91
Commit 5e644b91 authored Feb 09, 2023 by Woosuk Kwon
Hide whitespace changes
Inline Side-by-side

Showing with 20 additions and 0 deletions

cacheflow/utils.py cacheflow/utils.py +20 -0

No files found.
--- a/cacheflow/utils.py
+++ b/cacheflow/utils.py
+import enum
+
+
+class Device(enum.Enum):
+    GPU = enum.auto()
+    CPU = enum.auto()
+
+
+class Counter:
+
+    def __init__(self, start: int = 0) -> None:
+        self.counter = start
+
+    def next(self) -> int:
+        id = self.counter
+        self.counter += 1
+        return id
+
+    def reset(self) -> None:
+        self.counter = 0