unit test dataloader and distributedsampler

This commit is contained in:
lucasdelimanogueira 2024-05-25 11:55:59 -03:00
parent e590f57555
commit 10a88fe8ca
8 changed files with 204 additions and 13 deletions

View file

@ -33,7 +33,7 @@
"import norch\n",
"import norch.nn as nn\n",
"import norch.optim as optim\n",
"from norch.utils.data.dataloader import Dataloader\n",
"from norch.utils.data.dataloader import DataLoader\n",
"from norch.norchvision import transforms\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
@ -102,7 +102,7 @@
"device = \"cuda\" #cpu\n",
"epochs = 10\n",
"\n",
"transform = transforms.Sequential(\n",
"transform = transforms.Compose(\n",
" [\n",
" transforms.ToTensor(),\n",
" transforms.Reshape([-1, 784, 1])\n",
@ -116,7 +116,7 @@
")\n",
"\n",
"train_data, test_data = norch.norchvision.datasets.MNIST.splits(transform=transform, target_transform=target_transform)\n",
"train_loader = Dataloader(train_data, batch_size = BATCH_SIZE)\n",
"train_loader = DataLoader(train_data, batch_size = BATCH_SIZE)\n",
"\n",
"class MyModel(nn.Module):\n",
" def __init__(self):\n",