search

uber-research / UPSNet

UPSNet: A Unified Panoptic Segmentation Network
Other
649 stars 120 forks source link

A little bug in test script #83

Open pkuCactus opened 5 years ago

pkuCactus commented 5 years ago

Hi, when the length of dataset cannot be divided by the batch_size, then the origin code below https://github.com/uber-research/UPSNet/blob/aa8434e5a721ed217849607815304f68dfd7720a/upsnet/upsnet_end2end_test.py#L236 would cause cannot pin 'torch.cuda.FloatTensor' only CPU memory can be pinned and should change to data[k] = v.cpu().pin_memory().to(gpu_id, non_blocking=True) if torch.is_tensor(v) else v

rlangefe commented 3 years ago

I hit this too when trying to run. A fix for this would be helpful!