fix static cache

This commit is contained in:
2023-06-30 23:27:53 +02:00
parent 25d9814b70
commit d3ff1c05e1
+3 -3
View File
@@ -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