Use TaskNode Enum consistently

[kavyasrinet]

Description

As the title says.
Fixes #1094

Type of change

Please check the options that are relevant.

• Bug fix (non-breaking change that fixes an issue)
• Proposes a change (non-breaking change that isn't necessarily a bug)
• Refactor
• New feature (non-breaking change that adds a new functionality)
• Breaking change (fix or feature that would break some existing functionality downstream)
• This is a unit test
• Documentation only change
• Datasets Release
• Models Release

Type of requested review

• I want a thorough review of the implementation.
• I want a high level review.
• I want a deep design review.

Checklist:

• I have performed manual end-to-end testing of the feature in my environment.
• I have added Docstrings and comments to the code.
• I have made changes to existing documentation where needed.
• I have added tests that show that the PR is functional.
• New and existing unit tests pass locally with my changes.
• I have added relevant collaborators to review the PR before merge.
• [Polymetis only] I ran on hardware (1) all scripts in tests/scripts, (2) asv benchmarks.