diff options
author | InvincibleDude <81354513+InvincibleDude@users.noreply.github.com> | 2023-03-14 16:55:59 +0300 |
---|---|---|
committer | GitHub <noreply@github.com> | 2023-03-14 16:55:59 +0300 |
commit | f5e44364535ccc7efef445dacb6395c3942f2f17 (patch) | |
tree | ca7384e6225fccdae1e89db1e0ccc44dd60c7311 /modules/memmon.py | |
parent | f6e27378404631d951656388fc178b784fe1495b (diff) | |
parent | a9fed7c364061ae6efb37f797b6b522cb3cf7aa2 (diff) |
Merge branch 'master' into improved-hr-conflict-test
Diffstat (limited to 'modules/memmon.py')
-rw-r--r-- | modules/memmon.py | 12 |
1 files changed, 8 insertions, 4 deletions
diff --git a/modules/memmon.py b/modules/memmon.py index a7060f58..4018edcc 100644 --- a/modules/memmon.py +++ b/modules/memmon.py @@ -23,12 +23,16 @@ class MemUsageMonitor(threading.Thread): self.data = defaultdict(int) try: - torch.cuda.mem_get_info() + self.cuda_mem_get_info() torch.cuda.memory_stats(self.device) except Exception as e: # AMD or whatever print(f"Warning: caught exception '{e}', memory monitor disabled") self.disabled = True + def cuda_mem_get_info(self): + index = self.device.index if self.device.index is not None else torch.cuda.current_device() + return torch.cuda.mem_get_info(index) + def run(self): if self.disabled: return @@ -43,10 +47,10 @@ class MemUsageMonitor(threading.Thread): self.run_flag.clear() continue - self.data["min_free"] = torch.cuda.mem_get_info()[0] + self.data["min_free"] = self.cuda_mem_get_info()[0] while self.run_flag.is_set(): - free, total = torch.cuda.mem_get_info() # calling with self.device errors, torch bug? + free, total = self.cuda_mem_get_info() self.data["min_free"] = min(self.data["min_free"], free) time.sleep(1 / self.opts.memmon_poll_rate) @@ -70,7 +74,7 @@ class MemUsageMonitor(threading.Thread): def read(self): if not self.disabled: - free, total = torch.cuda.mem_get_info() + free, total = self.cuda_mem_get_info() self.data["free"] = free self.data["total"] = total |