fix training

This commit is contained in:
Andras Schmelczer 2024-04-13 15:54:34 +01:00
commit 4349f5af41
No known key found for this signature in database
GPG key ID: FC8F2C3D3D1A718C
4 changed files with 85068 additions and 118911 deletions

View file

@ -61,13 +61,14 @@ class HistogramDataset(Dataset):
edited = random_edit(original, seed=idx)
original_histogram = compute_histogram(
original, bins=self._bin_count, normalize=True
)
edited_histogram = compute_histogram(
edited, bins=self._bin_count, normalize=True
)
original_histogram = compute_histogram(
original, bins=self._bin_count, normalize=True
)
result = (
torch.tensor(edited_histogram, dtype=torch.float).unsqueeze(0),
torch.tensor(original_histogram, dtype=torch.float).unsqueeze(0),