## Prepare the data
(dataset, normalize=True, expand_dims=True)
| 46 | |
| 47 | |
| 48 | def preprocess(dataset, normalize=True, expand_dims=True): |
| 49 | """ |
| 50 | ## Prepare the data |
| 51 | """ |
| 52 | x, y = dataset.tensors |
| 53 | |
| 54 | if normalize: |
| 55 | # Scale images to the [0, 1] range |
| 56 | x = x.type(torch.float32) / 255 |
| 57 | |
| 58 | if expand_dims: |
| 59 | # Make sure images have shape (1, 28, 28) |
| 60 | x = torch.unsqueeze(x, 1) |
| 61 | |
| 62 | return TensorDataset(x, y) |
| 63 | |
| 64 | |
| 65 | def read(data_dir, split): |
no outgoing calls