code: def validation_prob_vs_label(
      bc_trainer: pl.LightningModule,
      batch: rlt.BehavioralCloningModelInput,
      batch_idx: int,
  ):
      masked_logits = bc_trainer.bc_net(
          batch.state,
          batch.possible_actions_mask,
      )
      labels = batch.action
      probs = torch.nn.functional.softmax(masked_logits)
      assert torch.allclose(labels.double(), probs.double(), atol=1e-1)
      return
  
imports: import logging
  import unittest
  from enum import Enum
  import pytorch_lightning
  import torch
  from pytorch_lightning import seed_everything
  from reagent.core import types
  from reagent.models.dqn import FullyConnectedDQN
  from reagent.optimizer.union import classes, Optimizer__Union
  from reagent.training.behavioral_cloning_trainer import BehavioralCloningTrainer
  from torch.utils.data import DataLoader
label_desc: Use image node_id 4 for calling a global function with example usage: validation_prob_vs_label(bc_trainer, batch, batch_idx) without return types