update forward pass test

This commit is contained in:
Johannes Fischer
2021-07-23 11:16:08 +02:00
parent acd2b730a6
commit 8a4094037a

View File

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