From d3ff1c05e118e61a2065581283a0877ab02d25a8 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Fri, 30 Jun 2023 23:27:53 +0200 Subject: [PATCH] fix static cache --- data/dataset.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/data/dataset.py b/data/dataset.py index 03e24c3..98be66f 100644 --- a/data/dataset.py +++ b/data/dataset.py @@ -105,15 +105,15 @@ class StaticCache(): def __getitem__(self, frame) -> dict: - if frame not in self.store: + if (frame.id, frame.garment) not in self.store: item = self.read(frame.id, frame.garment) stored_data = { key: torch.from_numpy(value).to('cpu') for key, value in item.items()} for key, value in stored_data.items(): if key.startswith('uv_') and value.shape[-1] == 3: stored_data[key] = torch.permute(value, [2,0,1]) - self.store[frame] = stored_data + self.store[(frame.id, frame.garment)] = stored_data - data = { key: value.to(self.device) for key, value in self.store[frame].items()} + data = { key: value.to(self.device) for key, value in self.store[(frame.id, frame.garment)].items()} return data