Hey, did you manage to resolve this issue? I am having the same problem with the model not training.
For example, this is what I get 
```
Epoch: 128 | Iteration: 0 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 1 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 2 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 3 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 4 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 5 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 6 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 7 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 8 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22222
Epoch: 128 | Iteration: 9 | Classification loss: 2.30212 | Regression loss: 0.37198 | Running loss: 0.22757
Epoch: 128 | Iteration: 10 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22272
Epoch: 128 | Iteration: 11 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22272
Epoch: 128 | Iteration: 12 | Classification loss: 2.30212 | Regression loss: 0.29100 | Running loss: 0.22790
Epoch: 128 | Iteration: 13 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 14 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 15 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 16 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 17 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 18 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 19 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 20 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 21 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 22 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 23 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 24 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 25 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 26 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 27 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 28 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 29 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22790
Epoch: 128 | Iteration: 30 | Classification loss: 0.00000 | Regression loss: 0.00000 | Running loss: 0.22304