pythonのgenerate処理でsegmentation faultが出て、原因が不明

Question

お世話になっております。
表題の件でわからず、困っております。
下位ディレクトリのtest_A,test_Bから画像を読み出し比較していく処理ですが、
どの部分が違反しているか、
おわかりになる方いらっしゃいましたらよろしくお願いします

cuda = True if torch.cuda.is_available() else False

generator = GeneratorUNet(in_channels=3)

if cuda:
  generator = generator.cuda()

generator.load_state_dict(torch.load("saved_models/%s/generator_%s.pth" % (opt.dataset_name, opt.saved_name), map_location='cpu'))

transforms_ = [
    transforms.Resize((256, 256), Image.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]

test_A, test_B = make_datapath_list(opt, phase='test')

dataloader = DataLoader(
    ImageDataset(test_A, test_B, opt, transforms_=transforms_),
    batch_size=opt.batch_size,
    shuffle=True,
)

Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor

for i, batch in enumerate(dataloader):
  real_A = Variable(batch["A"].type(Tensor))
  real_B = Variable(batch["B"].type(Tensor))

------ ココでsegmentation fault ------
  fake_B = generator(real_A)
--------------------------------------

  img_sample = torch.cat((real_A.data, fake_B.data, real_B), 0)
  save_image(img_sample, "results/{}/{}.png".format(opt.dataset_name, i), nrow=3)