fixing imports for moved test

Summary:

Test Plan: python test.py

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 7d4cabbf2656e32be71cedc5a5a201823048266d
Pull Request resolved: https://github.com/pytorch-labs/ao/pull/6