OctoberChang / X-Transformer

X-Transformer: Taming Pretrained Transformers for eXtreme Multi-label Text Classification
BSD 3-Clause "New" or "Revised" License
135 stars 28 forks source link

single-gpu eval fixes #5

Closed simonlevine closed 4 years ago

simonlevine commented 4 years ago

Hi, when running single GPU evaluation, there is an issue with the assumption:

assume self.model hasattr module because torch.nn.DataParallel

I noticed torch.nn.DataParallel is initiated upstream in the correct case of >1 GPUs, but here no such allotment for a single gpu was present. This resulted in the "module" not existing. I've fixed this with a simple if/then (a bit verbose, up to you to keep it of course) and it is tested as working on a model of type "bert". Thanks!