import hyper, gaussian, util, logging
import torch, random
from torch.autograd import Variable
from torch import nn, optim
from tqdm import trange, tqdm
from tensorboardX import SummaryWriter

from torchsample.metrics import CategoricalAccuracy

import torch.optim as optim

import torchvision
import torchvision.transforms as transforms

import matplotlib.pyplot as plt
import numpy as np

torch.manual_seed(1)
logging.basicConfig(filename='run.log',level=logging.INFO)
LOG = logging.getLogger()

"""
MNIST Autoencoder Experiment

"""
w = SummaryWriter()

BATCH = 1024
SHAPE = (28, 28)
MIDDLE = (10, )
EPOCHS = 350

CUDA = True

TYPE = 'free-weights'

gaussian.PROPER_SAMPLING = False

normalize = transforms.Compose(
    [transforms.ToTensor()
        #,transforms.Normalize((0.1307,), (0.3081,))
     ])

train = torchvision.datasets.MNIST(root='./data', train=True,
                                        download=True, transform=normalize)

trainloader = torch.utils.data.DataLoader(train, batch_size=BATCH,
                                          shuffle=True, num_workers=2)

test = torchvision.datasets.MNIST(root='./data', train=False,
                                       download=True, transform=normalize)

testloader = torch.utils.data.DataLoader(test, batch_size=BATCH,
                                         shuffle=True, num_workers=2)

if TYPE == 'non-adaptive':
    model = nn.Sequential(
        gaussian.ParamASHLayer(SHAPE, MIDDLE, k=16, additional=8, has_bias=True),
        nn.Sigmoid(),
        gaussian.ParamASHLayer(MIDDLE, SHAPE, k=16, additional=8, has_bias=True),
        nn.Sigmoid())
elif TYPE == 'free-weights':
    model = nn.Sequential(
        gaussian.CASHLayer(SHAPE, MIDDLE, k=750, additional=8, has_bias=False),
        nn.Sigmoid(),
        gaussian.CASHLayer(MIDDLE, SHAPE, k=1500, additional=8, has_bias=False),
        nn.Sigmoid())

if CUDA:
    model.apply(lambda t : t.cuda())

## SIMPLE
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

step = 0

for epoch in range(EPOCHS):
    for i, data in tqdm(enumerate(trainloader, 0)):

        # get the inputs
        inputs, _ = data
        inputs = inputs.squeeze(1)

        if CUDA:
            inputs = inputs.cuda()

        # wrap them in Variables
        inputs, targets = Variable(inputs), Variable(inputs)

        # forward + backward + optimize
        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()

        optimizer.step()

        w.add_scalar('autoencoder/train-loss', loss.data[0], step)

        step += 1

    total = 0.0
    num = 0
    for i, data in tqdm(enumerate(testloader, 0)):
        # get the inputs
        inputs, _ = data
        inputs = inputs.squeeze(1)

        if CUDA:
            inputs = inputs.cuda()

        # wrap them in Variables
        inputs, targets = Variable(inputs), Variable(inputs)

        # forward + backward + optimize
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)

        total += loss.data
        loss = None

        num += 1

        if i < 1:
            inputs, outputs = inputs.unsqueeze(1), outputs.unsqueeze(1)

            plt.figure(figsize=(16, 4))
            plt.imshow(np.transpose(torchvision.utils.make_grid(inputs.data[:16,:]).cpu().numpy(), (1, 2, 0)), interpolation='nearest')
            plt.savefig('input.{}.pdf'.format(epoch))

            plt.figure(figsize=(16, 4))
            plt.imshow(np.transpose(torchvision.utils.make_grid(outputs.data[:16,:]).cpu().numpy(), (1, 2, 0)), interpolation='nearest')
            plt.savefig('output.{}.pdf'.format(epoch))

    epoch_loss = total/(num * BATCH)

    w.add_scalar('autoencoder/epoch-test-loss', epoch_loss.data[0], epoch)
    print('EPOCH {}: {} loss per instance '.format(epoch, epoch_loss.data[0]))

print('Finished Training.')


@@ -617,19 +617,20 @@
 %0ACUDA = 
-Tru
+Fals
 e%0A%0ATYPE 
@@ -1724,33 +1724,34 @@
 750, additional=
-8
+32
 , has_bias=False
@@ -1828,33 +1828,34 @@
 500, additional=
-8
+32
 , has_bias=False
@@ -2620,24 +2620,25 @@
  step += 1%0A%0A
+%0A
     total = 
@@ -3101,16 +3101,22 @@
 otal += 
+float(
 loss.dat
@@ -3112,24 +3112,28 @@
 at(loss.data
+%5B0%5D)
 %0A        los

Fix test-time memory bug (really this time)