Fix edge grad assert mismatch

Summary:
* The sizes for vi in the edge_grad_estimator_fwd assert were not updated after D68534639 expanded the dimension to 3
* This updates the size in the assert and adds an explicit call to edge_grad_estimator_fwd (a no-op) in the autograd implementation to make sure the sizes are checked

Reviewed By: HapeMask, phg1024

Differential Revision: D72433642

fbshipit-source-id: 49dd82e0a07fe174c2157b362eedf464984d386d
8 jobs for main
Status Job ID Name Coverage
  Build
canceled #19581
build

 
  Test
canceled #19584
allowed to fail
bandit-sast

canceled #19583
allowed to fail
code_quality

canceled #19585
allowed to fail
eslint-sast

canceled #19586
allowed to fail
flawfinder-sast

canceled #19588
allowed to fail
secret_detection

canceled #19587
allowed to fail
semgrep-sast

canceled #19582
test