This commit is contained in:
Arec
2021-07-23 04:44:57 -07:00

View File

@@ -10,8 +10,12 @@ with open(config_path, 'r') as cfg:
filestr = 'tests/policies/base_'
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])