I hope to add a pretrained network in transforms.lambda() function. and transform the input images into 2D tensors. Then, using the 2D tensors as the input to another network.
I meet a problem in DataLoader. When the transforms don’t compose the lambda transform, DataLoaderworks well. But I add the lambda transform, there is a error when I call:

next(iter(dset_loaders[‘val’]))

How should I solve for the problem?
Error

RuntimeError                              Traceback (most recent call last)
 in ()
----> 1 next(iter(dset_loaders[‘val’]))


/project/focus/hong/anaconda3/lib/python3.5/site-packages/torch/utils/data/dataloader.py in next(self)
172                 self.reorder_dict[idx] = batch
173                 continue
–> 174             return self._process_next_batch(batch)
175
176     next = next  # Python 2 compatibility


/project/focus/hong/anaconda3/lib/python3.5/site-packages/torch/utils/data/dataloader.py in _process_next_batch(self, batch)
196         self._put_indices()
197         if isinstance(batch, ExceptionWrapper):
–> 198             raise batch.exc_type(batch.exc_msg)
199         return batch
200
RuntimeError: Traceback (most recent call last):
File “/project/focus/hong/anaconda3/lib/python3.5/site-packages/torch/utils/data/dataloader.py”, line 34, in _worker_loop
samples = collate_fn([dataset[i] for i in batch_indices])
File “/project/focus/hong/anaconda3/lib/python3.5/site-packages/torch/utils/data/dataloader.py”, line 34, in 
samples = collate_fn([dataset[i] for i in batch_indices])
File “/project/focus/hong/anaconda3/lib/python3.5/site-packages/torchvision-0.1.8-py3.5.egg/torchvision/datasets/folder.py”, line 67, in getitem
img = self.transform(img)
File “/project/focus/hong/anaconda3/lib/python3.5/site-packages/torchvision-0.1.8-py3.5.egg/torchvision/transforms.py”, line 29, in call
img = t(img)
File “/project/focus/hong/anaconda3/lib/python3.5/site-packages/torchvision-0.1.8-py3.5.egg/torchvision/transforms.py”, line 184, in call
return self.lambd(img)
File “”, line 4, in 
transforms.Lambda(lambda x: FM(resnet, x))])
File “”, line 23, in FM
input_imgs = Variable(imgs.cuda())
File “/project/focus/hong/anaconda3/lib/python3.5/site-packages/torch/_utils.py”, line 65, in cuda
return new_type(self.size()).copy(self, async)
RuntimeError: cuda runtime error (3) : initialization error at /py/conda-bld/pytorch_1490983232023/work/torch/lib/THC/generic/THCStorage.c:55


#My code is attached below

Network function for transforms
def FM(model, img):
[d,h,w] = img.size()
if w > h:
    w_new, h_new, s_new = int((w - h)/2), 0, h // NGSIZE
else:
    w_new, h_new, s_new = 0, int((h - w)/2), w // NGSIZE



imgs = torch.Tensor(NGSIZE**2,d,s_new,s_new)



# crop the center part of the image
i = 0
for nh in range(NGSIZE):
    for nw in range(NGSIZE):
        Hsta = h_new + nh * s_new
        Wsta = w_new + nw * s_new
        Hend = h_new + nh * s_new + s_new
        Wend = w_new + nw * s_new + s_new
        imgs[i,:,:,:] = img[:,Hsta:Hend,Wsta:Wend]
        i += 1
        
if use_gpu:
    input_imgs = Variable(imgs.cuda())
else:
    input_imgs = Variable(imgs)



# forward
outputs = model(input_imgs).squeeze()

return outputs.cpu().data.view(3,3,512)



import resnet
resnet_ptra = models.resnet18(pretrained=True)
resnet_ptra.train(False)
resnet = nn.Sequential(*list(resnet_ptra.children())[:-1],nn.AvgPool2d(3))
for param in resnet.parameters(): param.requires_grad = False
if use_gpu:
resnet = resnet.cuda()
data pre-processing
data_transforms = transforms.Compose([transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
transforms.Lambda(lambda x: FM(resnet, x))])


data loader
dsets = {x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms) for x in [‘train’, ‘val’]}


dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=4,
shuffle=False, num_workers=8)
for x in [‘train’, ‘val’]}