OIM loss initialize error

[yaodonggggggg]

I test the examples/oim_loss.py with python examples/oim_loss.py -d cuhk01 -b 2 -j 0 -a resnet50 --logs-dir logs/softmax-loss/cuhk01-resnet50. And I got the model.module.num_features 2048 ,but when I run the resnet.py in your code with

model = create('resnet50', num_features=0, norm=True,dropout=0,num_classes=4)
x = torch.rand(2, 3, 256,128)
out = model(x)
print("out.shape",out.shape)
loss = OIMLoss(2048, 4)
target = torch.empty(2, dtype=torch.long).random_(4)
print(target)
los, output = loss(out, target)

But got

File "F:/python_workspace/open-reid-master/reid/loss/oim.py", line 16, in forward
outputs = inputs.mm(self.lut.t())
RuntimeError: size mismatch, m1: [2 x 4], m2: [2048 x 4]

How to deal with the problem?