Training of classifier breaks when casting uint8 to float

Hello guys,
I was just playing around with a toy example on MNIST, but i was using a conv1D network. Achieved around 96% after 5 epochs. Then i wanted to implement my own normalization function, but had to cast the data to float() first because you cant use mean() on a bytetensor. Just calling .float() on the data caused the network to suddenly drop to 50% for some odd reason. Does anyone know what might be happening here? I have the 2 notebooks on collab:

https://drive.google.com/open?id=1nNzMyf--R-j8K4GxUSHwArF2nzvinxVi
https://drive.google.com/open?id=1PAZ4dCxKD-pR8LRSnGQgbCFDJ143nQRI


EPOCHS = 5
BS = 128
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
transform = transforms.Compose([transforms.ToTensor()])
mnist_train = torchvision.datasets.MNIST('.', download=True, transform=transform,train=True)
mnist_test = torchvision.datasets.MNIST('.', download=True, transform=transform, train=False)
#####
#mnist_train.data = mnist_train.data.float()
#mnist_test.data = mnist_test.data.float()

trainloader = torch.utils.data.DataLoader(mnist_train,batch_size=BS,shuffle=True)
testloader = torch.utils.data.DataLoader(mnist_test,batch_size=10000,shuffle=False)

class conv_net(nn.Module):
    def __init__(self, lr):
        super(conv_net, self).__init__()
        self.layer_1 = nn.Conv1d(1, 32, kernel_size=7, stride=5)
        self.layer_2 = nn.Conv1d(32, 32, kernel_size=7, stride=5)
        self.out = nn.Linear(960, 10)
        self.optimizer = optim.Adam(self.parameters(), lr = lr)
    def forward(self, x):
        #print(x.dtype)
        x = x.to(device)
        x = x.view((x.size()[0], 1, -1))
        x = F.leaky_relu(self.layer_1(x), negative_slope=0.1)
        x = F.relu(self.layer_2(x)).view(x.size(0),-1)
        x = self.out(x)
        return x

model = conv_net(1e-3).to(device)
criterion = torch.nn.CrossEntropyLoss()

for epoch in range(EPOCHS):
    for x_train, y_train in iter(trainloader):
        model.train()
        y_train = y_train.to(device)
        pred = model(x_train)
        loss = criterion(pred, y_train)
        loss.backward()
        model.optimizer.step()
        model.optimizer.zero_grad()
       
    with torch.no_grad():
        counter = 0
        for x_test, y_test in iter(testloader):
            y_test = y_test.to(device)
            model.eval()
            pred = model(x_test)
            pred = torch.argmax(F.log_softmax(pred), dim = 1)
            correct = (pred==y_test).sum().item()
        print(f'Test - Accuracy: {(correct/len(y_test)):.4f}') for epoch in range(EPOCHS):
    for x_train, y_train in iter(trainloader):
        model.train()
        y_train = y_train.to(device)
        pred = model(x_train)
        loss = criterion(pred, y_train)
        loss.backward()
        model.optimizer.step()
        model.optimizer.zero_grad()
       
    with torch.no_grad():
        counter = 0
        for x_test, y_test in iter(testloader):
            y_test = y_test.to(device)
            model.eval()
            pred = model(x_test)
            pred = torch.argmax(F.log_softmax(pred), dim = 1)
            correct = (pred==y_test).sum().item()
        print(f'Test - Accuracy: {(correct/len(y_test)):.4f}')

Seems like the dataloader is gonna cast it to float anyways before serving the model. So basically the model gets the exact same data for both runs? Cant seem to be related to initalization, i tried both runs more than 20 times. Same result every time

The Image.fromarray operation fails, if you manipulate the underlying data in this line of code.

Example:

dataset = datasets.MNIST(
    root='PATH',
    transform=transforms.ToTensor(),
    download=True)

img = dataset.data[0]
img = Image.fromarray(img.numpy(), mode='L')
img.show() # works

img = dataset.data[0].float()
img = Image.fromarray(img.numpy(), mode='L')
img.show() # interleaved output