Spaces:

hugging-apps
/

codecsplat

Running on Zero

App Files Files Community

codecsplat / src /misc /step_tracker.py

multimodalart HF Staff

CodecSplat demo

7405e85 verified 8 days ago

Raw History Blame Contribute Delete

506 Bytes

from multiprocessing import RLock

import torch
from jaxtyping import Int64
from torch import Tensor


class StepTracker:
    lock: RLock
    step: Int64[Tensor, ""]

    def __init__(self):
        self.lock = torch.multiprocessing.RLock()
        self.step = torch.tensor(0, dtype=torch.int64).share_memory_()

    def set_step(self, step: int) -> None:
        with self.lock:
            self.step.fill_(step)

    def get_step(self) -> int:
        with self.lock:
            return self.step.item()

	from multiprocessing import RLock

	import torch
	from jaxtyping import Int64
	from torch import Tensor


	class StepTracker:
	lock: RLock
	step: Int64[Tensor, ""]

	def __init__(self):
	self.lock = torch.multiprocessing.RLock()
	self.step = torch.tensor(0, dtype=torch.int64).share_memory_()

	def set_step(self, step: int) -> None:
	with self.lock:
	self.step.fill_(step)

	def get_step(self) -> int:
	with self.lock:
	return self.step.item()