text: #
  # Copyright 1993-2018 NVIDIA Corporation.  All rights reserved.
  #
  # NOTICE TO LICENSEE:
  #
  # This source code and/or documentation ("Licensed Deliverables") are
  # subject to NVIDIA intellectual property rights under U.S. and
  # international Copyright laws.
  #
  # These Licensed Deliverables contained herein is PROPRIETARY and
  # CONFIDENTIAL to NVIDIA and is being provided under the terms and
  # conditions of a form of NVIDIA software license agreement by and
  # between NVIDIA and Licensee ("License Agreement") or electronically
  # accepted by Licensee.  Notwithstanding any terms or conditions to
  # the contrary in the License Agreement, reproduction or disclosure
  # of the Licensed Deliverables to any third party without the express
  # written consent of NVIDIA is prohibited.
  #
  # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
  # LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE
  # SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE.  IT IS
  # PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND.
  # NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED
  # DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
  # NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
  # NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
  # LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY
  # SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
  # DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
  # WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
  # ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
  # OF THESE LICENSED DELIVERABLES.
  #
  # U.S. Government End Users.  These Licensed Deliverables are a
  # "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT
  # 1995), consisting of "commercial computer software" and "commercial
  # computer software documentation" as such terms are used in 48
  # C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government
  # only as a commercial end item.  Consistent with 48 C.F.R.12.212 and
  # 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all
  # U.S. Government End Users acquire the Licensed Deliverables with
  # only those rights set forth herein.
  #
  # Any use of the Licensed Deliverables in individual and commercial
  # software must include, in the user documentation and internal
  # comments to the code, the above Disclaimer and U.S. Government End
  # Users Notice.
  #
  
  try:
      import torch
      import torch.nn as nn
      import torch.nn.functional as F
      import torch.optim as optim
      from torchvision import datasets, transforms
      from torch.autograd import Variable
  except ImportError as err:
      raise ImportError("""ERROR: Failed to import module ({})
  Please make sure you have PyTorch installed.
  For installation instructions, see:
  http://pytorch.org/""".format(err))
  
  import numpy as np
  import os
  
  BATCH_SIZE = 64
  TEST_BATCH_SIZE = 100
  EPOCHS = 10
  LEARNING_RATE = 0.01
  SGD_MOMENTUM = 0.9
  SEED = 1
  LOG_INTERVAL = 100
  
  #Enable Cuda
  torch.cuda.manual_seed(SEED)
  
  #Dataloader
  kwargs = {'num_workers': 1, 'pin_memory': True}
  train_loader  = torch.utils.data.DataLoader(
      datasets.MNIST('/tmp/mnist/data', train=True, download=True, transform=transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize((0.1307,), (0.3081,))
          ])),
      batch_size=BATCH_SIZE,
      shuffle=True,
      **kwargs)
  
  test_loader = torch.utils.data.DataLoader(
      datasets.MNIST('/tmp/mnist/data', train=False, transform=transforms.Compose([
          transforms.ToTensor(),
          transforms.Normalize((0.1307,), (0.3081,))
          ])),
      batch_size=TEST_BATCH_SIZE,
      shuffle=True,
      **kwargs
      )
  
  #Network
  class Net(nn.Module):
      def __init__(self):
          super(Net, self).__init__()
          self.conv1 = nn.Conv2d(1, 20, kernel_size=5)
          self.conv2 = nn.Conv2d(20, 50, kernel_size=5)
          self.conv2_drop = nn.Dropout2d()
          self.fc1 = nn.Linear(800, 500)
          self.fc2 = nn.Linear(500, 10)
  
      def forward(self, x):
          x = F.max_pool2d(self.conv1(x), kernel_size=2, stride=2)
          x = F.max_pool2d(self.conv2(x), kernel_size=2, stride=2)
          x = x.view(-1, 800)
          x = F.relu(self.fc1(x))
          x = self.fc2(x)
          return F.log_softmax(x, dim=1)
  
  model = Net()
  model.cuda()#.half()
  
  optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=SGD_MOMENTUM)
  
  def train(epoch):
      model.train()
      for batch, (data, target) in enumerate(train_loader):
          data, target = data.cuda(), target.cuda()
          data, target = Variable(data), Variable(target)
          optimizer.zero_grad()
          output = model(data)
          loss = F.nll_loss(output, target)
          loss.backward()
          optimizer.step()
          if batch % LOG_INTERVAL == 0:
              print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(epoch, batch * len(data), len(train_loader.dataset), 100. * batch / len(train_loader), loss.data.item()))
  
  def test(epoch):
      model.eval()
      test_loss = 0
      correct = 0
      for data, target in test_loader:
          data, target = data.cuda(), target.cuda()
          with torch.no_grad():
              data, target = Variable(data), Variable(target)
          output = model(data)
          test_loss += F.nll_loss(output, target).data.item()
          pred = output.data.max(1)[1]
          correct += pred.eq(target.data).cpu().sum()
      test_loss /= len(test_loader)
      print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)))
  
  def learn():
      for e in range(EPOCHS):
          train(e + 1)
          test(e + 1)
  
  def get_trained_model():
      return model.state_dict()
  
  def get_testcase():
      return next(iter(test_loader))
  
  if __name__ == '__main__':
      learn()
      print(model)
      torch.save(model.state_dict(), os.path.dirname(os.path.realpath(__file__)) + "/trained_mnist.pyt")
  