I finally had time to create a minimal example. Both train and validation loss goes to nan pretty quickly.
import time
import numpy as np
import argparse

import torch
import torchvision
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as tforms

import torchvision as vsn

parser = argparse.ArgumentParser(description='CIFAR10 Training')
parser.add_argument('--lr', default=0.001, type=float,
                    help='initial learning rate')
parser.add_argument('--epochs', default=20, type=int, 
                    help='number of epochs')
parser.add_argument('--batch_size', default=256, type=int, 
                    help='size of batches')
parser.add_argument('--gpu', default=0, type=int, 
                    help='which gpu to run')
args = parser.parse_args()

class ResNet(nn.Module):
    def __init__(self):
        super(ResNet, self).__init__()
        self.resnet_layers = list(vsn.models.resnet18(pretrained=False).children())
        self.maxpool = self.resnet_layers[3]
        # encoding layers
        self.encode_a = nn.Sequential(*self.resnet_layers[:3])
        self.encode_b = self.resnet_layers[4]
        self.encode_c = self.resnet_layers[5]
        self.encode_d = self.resnet_layers[6]
        self.encode_e = self.resnet_layers[7]
        # avgpool
        self.avgpool = nn.AvgPool2d(kernel_size=2)
        # output layer
        self.linear = nn.Linear(512, 10)
        # dropout
        self.dropout = nn.Dropout2d(p=0.15)

    def forward(self, x):
        x = self.encode_a(x)
        x = self.encode_b(x)
        x = self.encode_c(x)
        x = self.encode_d(x)
        x = self.encode_e(x)
        x = self.avgpool(x)
        out = self.linear(x.squeeze())
        
        return out


if args.gpu == 99:
	device = torch.device("cuda:0")
else:
	torch.cuda.set_device(args.gpu)
	device = torch.device('cuda:{}'.format(args.gpu) if torch.cuda.is_available() else 'cpu')

transforms = tforms.Compose([tforms.RandomHorizontalFlip(p=0.5),
	                         tforms.RandomCrop(size=32, padding=4),
	                         tforms.ToTensor(),
	                         tforms.Normalize((0.5, 0.5, 0.5),
	                             	          (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='../data', train=True,
                                        download=True, transform=transforms)

trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='../data', train=False,
                                       download=True, transform=tforms.Compose([tforms.ToTensor(),
                                                                                tforms.Normalize(
                                                                                    (0.5,0.5,0.5),
                                                                                    (0.5,0.5,0.5))]))
testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size,
                                         shuffle=False, num_workers=2)


print('Using {} labeled  data'.format(len(trainset)))
print('Using {} validation data'.format(len(testset)))

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

net = ResNet()
if args.gpu == 99:
     net = nn.DataParallel(net)
net.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=args.lr)

def train():
    net.train(True)
    running_loss = 0.

    for i, data in enumerate(trainloader):
        imgs, labels = data

        optimizer.zero_grad()

        preds = net(imgs.to(device))
        loss = criterion(preds, labels.to(device))
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

    avg_loss = running_loss / (len(trainloader.dataset) / args.batch_size)
    print('Avg Train Loss: {:.4}'.format(avg_loss))
    return avg_loss

def valid():
    net.eval()

    running_loss = 0.
    running_corr = 0.
    total = 0.

    with torch.no_grad():
        for data in testloader:
            imgs, labels = data

            preds = net(imgs.to(device))

            _, predicted = torch.max(preds.data, 1)
            total += labels.size(0)
            running_corr += (predicted == labels.to(device)).sum().item()

            running_loss += criterion(preds, labels.to(device)).item()

    avg_loss = running_loss / (len(testloader.dataset) / args.batch_size)
    acc = running_corr / total
    print('Avg Valid Loss: {:.4}, Valid Acc: {:.4}'.format(avg_loss, acc))
    return avg_loss, acc

train_losses = []
valid_losses = []
valid_accs = []
try:
    for e in range(args.epochs):
        start = time.time()
        print('\n' + 'Epoch {}/{}'.format(e, args.epochs))
        train_loss = train()
        valid_loss, valid_acc = valid()

        train_losses.append(train_loss)
        valid_losses.append(valid_loss)
        valid_accs.append(valid_acc)
        end = time.time() - start
        print('{} seconds'.format(end))

except KeyboardInterrupt:
    pass

import pandas as pd

out_dict = {'train_losses': train_losses, 
            'valid_losses': valid_losses,
            'valid_accs': valid_accs}

out_log = pd.DataFrame(out_dict)
print(out_log.head())

out_log.to_csv('cifar10_resnet_gpu-{}.csv'.format(args.gpu), index=False)