import torch
from torchvision import datasets, transforms
import numpy as np

from deeprobust.image.defense.trades import TRADES
from deeprobust.image.netmodels.CNN import Net

train_loader = torch.utils.data.DataLoader(
                datasets.MNIST('deeprobust/image/defense/data', train = True, download = True,
                transform = transforms.Compose([transforms.ToTensor()])),
                batch_size = 100,
                shuffle = True)  

test_loader = torch.utils.data.DataLoader(
            datasets.MNIST('deeprobust/image/defense/data', train = False,
            transform = transforms.Compose([transforms.ToTensor()])),
            batch_size = 1000,
            shuffle = True)  


model = Net()
defense = TRADES(model,'cuda')
defense.generate(train_loader, test_loader)