diff options
Diffstat (limited to 'experiments/bci_smoke.py')
| -rw-r--r-- | experiments/bci_smoke.py | 9 |
1 files changed, 8 insertions, 1 deletions
diff --git a/experiments/bci_smoke.py b/experiments/bci_smoke.py index 4724d26..1fda89e 100644 --- a/experiments/bci_smoke.py +++ b/experiments/bci_smoke.py @@ -104,7 +104,14 @@ def check_temporal_difference_role_vectorizer(): cfg.episodes_per_day, cfg.n_neurons, generator=generator) xi = torch.empty_like(soma).bernoulli_( 0.5, generator=generator).mul_(2).sub_(1) - target = model.causal_role_targets(soma, xi) + cursor_plus = (soma + cfg.perturb_sigma * xi) @ model.role + cursor_minus = (soma - cfg.perturb_sigma * xi) @ model.role + target = model.causal_role_targets(cursor_plus, cursor_minus, xi) + saved_role = model.role.clone() + model.role.zero_() + assert torch.equal( + target, model.causal_role_targets(cursor_plus, cursor_minus, xi)) + model.role.copy_(saved_role) mean_target = target.mean(0) cosine = torch.nn.functional.cosine_similarity( mean_target, model.role, dim=0).item() |