[sparsity] Fix and enable the pruning tests (#66411)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/66411
The original tests were disabled, and had some bugs. This fixes those unittests.
Test Plan: Imported from OSS
Reviewed By: HDCharles
Differential Revision: D31590678
Pulled By: z-a-f
fbshipit-source-id: ddbed34cc01d5f15580cb8f0033416f2f9780068