1
0
Fork 0

add subscript for fake image

This commit is contained in:
Sean Sube 2023-01-16 12:02:37 -06:00
parent 402fa4f687
commit 0f2d6d2ec7
1 changed files with 3 additions and 0 deletions

View File

@ -29,6 +29,9 @@ class ONNXImage():
self.source = source self.source = source
self.data = self self.data = self
def __getitem__(self, *args):
return torch.from_numpy(self.source.__getitem__(*args)).to(torch.float32)
def squeeze(self): def squeeze(self):
self.source = np.squeeze(self.source, (0)) self.source = np.squeeze(self.source, (0))
return self return self