fix static cache
This commit is contained in:
+3
-3
@@ -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
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user