liutinglt / CE2P

214 stars 41 forks source link

why don't push images in GPU when training #25

Closed Anikily closed 5 years ago

Anikily commented 5 years ago

            images, labels, edges, _ = batch
            labels = labels.long().cuda(non_blocking=True)
            edges = edges.long().cuda(non_blocking=True)

            preds = model(images)

images still be the torch.float before put into model.