read rgb image as well

This commit is contained in:
2023-01-06 11:46:57 +00:00
parent 3c2c851ba5
commit 304a66e5b7
2 changed files with 12 additions and 12 deletions
+2 -3
View File
@@ -57,9 +57,8 @@ def batch_to_device(batch, device):
if __name__ == '__main__':
device = "cuda" if torch.cuda.is_available() else "cpu"
train_dataset = Dataset('/home/workspace/local-data/train')
valid_dataset = Dataset('/home/workspace/local-data/valid')
train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_all.txt')
train_dataloader = DataLoader(train_dataset, batch_size=4, drop_last=True, num_workers=2)
valid_dataloader = DataLoader(valid_dataset, batch_size=4, drop_last=True, num_workers=2)
train_accumulator = Accumulator()