Temporarily fixed a bug for vmap

[fantaosha]

Motivation and Context

How Has This Been Tested

Types of changes

• Docs change / refactoring / dependency upgrade
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)

Checklist

• My code follows the code style of this project.
• My change requires a change to the documentation.
• I have updated the documentation accordingly.
• I have read the CONTRIBUTING document.
• I have completed my CLA (see CONTRIBUTING)
• I have added tests to cover my changes.
• All new and existing tests passed.

[luisenp]

It does seem that there is some bug in torch when autograd generates a zero gradient for an operation. I passed a simple repro to @zou3519 to double check this.

That being said, it seems that the reason we get a zero gradient is because the jacobian autograd functions return poses for selected links + ancestors (see here and here), but only the poses of links with parents are used (all but poses[0], if I understand correctly).

Is there a reason why we can't make selected_links only contain the links whose poses will be used by jacobian computation? (in line 138). As far as I can tell, the first element of selected_poses will never be used (it's not used by the backward_helper and it's not used by the loop in line 155). Is this correct?