Validation for weights tensor to have same number of classes as the number of annotated classes

In the test data that was annotate by Rama, we had 4 classes annotated, and when we tried to train the model, the weights from a previous training run were used, which had 2 values due to 2 classes. There was no validation done to match the new number of annotated classes, but the training fails. A check needs to placed to match the number of classes and weights.