fix a bug in single GPU version, convert weight shape to 2d