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