Error when importing mnist1d.data and mnist1d.utils in Colab notebook

When trying to use the repo in a Colab notebook, in the import section:

# Setup
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import datasets, transforms

import math
from torchsummary import summary
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt

# Additional Setup for MNIST-1D
!git clone https://github.com/greydanus/mnist1d
import mnist1d
from mnist1d.data import get_templates, get_dataset_args, get_dataset
from mnist1d.utils import set_seed, plot_signals, ObjectView, from_pickle

# Additional Setup to use Tensorboard
!pip install -q tensorflow
%load_ext tensorboard

I obtain the following error:

Cloning into 'mnist1d'...
remote: Enumerating objects: 537, done.
remote: Counting objects: 100% (156/156), done.
remote: Compressing objects: 100% (86/86), done.
remote: Total 537 (delta 96), reused 108 (delta 67), pack-reused 381 (from 1)
Receiving objects: 100% (537/537), 14.92 MiB | 15.39 MiB/s, done.
Resolving deltas: 100% (249/249), done.
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
[<ipython-input-1-33e438737ccd>](https://localhost:8080/#) in <cell line: 19>()
     17 get_ipython().system('git clone https://github.com/greydanus/mnist1d')
     18 import mnist1d
---> 19 from mnist1d.data import get_templates, get_dataset_args, get_dataset
     20 from mnist1d.utils import set_seed, plot_signals, ObjectView, from_pickle
     21 

ModuleNotFoundError: No module named 'mnist1d.data'

I solved it by editing the data.py and utils.py files on lines 8, 9 (data.py) and 8 (utilis.py):

data.py: modified

from mnist1d.transform import transform
from mnist1d.utils import from_pickle, to_pickle, ObjectView, set_seed

with

from .transform import transform
from .utils we import from_pickle, to_pickle, ObjectView, set_seed

utils.py modified from mnist1d.transform import transform with from .transform import transform

greydanus / mnist1d

Error when importing mnist1d.data and mnist1d.utils in Colab notebook #16