Fix skipped tests in test_model_parallel.py (#3203)

Summary:
Pull Request resolved: #3203

Fix skipped tests. See more context in D78355780.

https://www.internalfb.com/intern/test/844425010918803?ref_report_id=0

Reviewed By: jeffkbkim

Differential Revision: D78461387

fbshipit-source-id: cfb043f4c15583006ce0dd0d3ea2038d4547be4c

Author: jd7-tr

torchrec/distributed/test_utils/test_model_parallel.py

@@ -467,11 +467,10 @@ def test_sharding_cw(
         data_type: DataType,
         allow_zero_batch_size: bool,
     ) -> None:
-        if (
-            self.device == torch.device("cpu")
-            and kernel_type != EmbeddingComputeKernel.FUSED.value
-        ):
-            self.skipTest("CPU does not support uvm.")
+        assume(
+            self.device != torch.device("cpu")
+            or kernel_type == EmbeddingComputeKernel.FUSED.value
+        )
 
         sharding_type = ShardingType.COLUMN_WISE.value
         assume(
@@ -548,11 +547,10 @@ def test_sharding_twcw(
         variable_batch_size: bool,
         data_type: DataType,
     ) -> None:
-        if (
-            self.device == torch.device("cpu")
-            and kernel_type != EmbeddingComputeKernel.FUSED.value
-        ):
-            self.skipTest("CPU does not support uvm.")
+        assume(
+            self.device != torch.device("cpu")
+            or kernel_type == EmbeddingComputeKernel.FUSED.value
+        )
 
         sharding_type = ShardingType.TABLE_COLUMN_WISE.value
         assume(
@@ -629,11 +627,10 @@ def test_sharding_tw(
         variable_batch_size: bool,
         data_type: DataType,
     ) -> None:
-        if (
-            self.device == torch.device("cpu")
-            and kernel_type != EmbeddingComputeKernel.FUSED.value
-        ):
-            self.skipTest("CPU does not support uvm.")
+        assume(
+            self.device != torch.device("cpu")
+            or kernel_type == EmbeddingComputeKernel.FUSED.value
+        )
 
         sharding_type = ShardingType.TABLE_WISE.value
         assume(