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