dataset fix

This commit is contained in:
2023-04-07 11:01:30 +00:00
parent b4be8bf80d
commit 1ad1fbfe1e
+2 -2
View File
@@ -141,7 +141,7 @@ class Dataset(Dataset):
if self.cache is not None: if self.cache is not None:
if frame in self.cache.keys(): if frame in self.cache.keys():
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()} data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in self.cache[frame].items()}
data = { **item, **self.static_cache[frame.id] } data = { **data, **self.static_cache[frame.id] }
return data return data
separator = '_' separator = '_'
@@ -213,7 +213,7 @@ class Dataset(Dataset):
self.cache[frame] = data self.cache[frame] = data
data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()} data = { key: (value.to(self.device) if type(value) == torch.Tensor else value) for key, value in data.items()}
data = { **item, **self.static_cache[frame.id] } data = { **data, **self.static_cache[frame.id] }
return data return data