update forward pass test

tests/policies/test_forward_pass.py

@@ -10,8 +10,12 @@ with open(config_path, 'r') as cfg:
 filestr = 'tests/policies/base_'
 
 
-batch = pickle.load(open(filestr+'_test_batch.pkl', 'rb'))
+def test_forward_pass():
+batch = pickle.load(open(filestr + '_test_batch.pkl', 'rb'))
 batch['relative_state'] = batch['relative_state'].float()
 policy = BehaviorCloningPolicy.load_model(config, filestr)
 policy.policy = policy.policy.type(batch['state'].dtype)
 action = policy(batch)
+i = torch.where(action.isnan())[0]
+for key in batch.keys():
+    print(batch[key][i])