code: def num_classes_from_teacher(teacher: BaseImageEncoder):
      teacher_classifier = teacher.classifier
      if isinstance(teacher_classifier, (torch.nn.Linear, LinearLayer)):
          return teacher_classifier.out_features
      elif isinstance(teacher_classifier, torch.nn.Sequential):
          last_layer = teacher_classifier[-1]
          assert isinstance(last_layer, (torch.nn.Linear, LinearLayer))
          return last_layer.out_features
      else:
          raise NotImplementedError
  
file_path: ml-cvnets/tests/loss_fns/test_distillation_loss.py
imports: import pytest
  import torch
  from cvnets.layers import LinearLayer
  from cvnets.models.classification.base_image_encoder import BaseImageEncoder
  from loss_fn.distillation.hard_distillation import HardDistillationLoss
  from loss_fn.distillation.soft_kl_distillation import SoftKLLoss
  from tests.configs import get_config
  from tests.test_utils import unset_pretrained_models_from_opts
label_desc: Use image node_id 3 for calling a global function with example usage: num_classes_from_teacher(teacher) and returns: teacher_classifier, last_layer