[pytorch] Fixes unit test failure on GPU.

[frankfliu]

Change-Id: Ia2d1271ad97ccc2f59059981a7534e2470160943

Description

Brief description of what this PR is about

• If this change is a backward incompatible change, why must this change be made?
• Interesting edge cases to note here

[codecov-commenter]

Codecov Report

Merging #1430 (ad6efd3) into master (bb5073f) will decrease coverage by 1.41%.
The diff coverage is 47.78%.

@@             Coverage Diff              @@
##             master    #1430      +/-   ##
============================================
- Coverage     72.08%   70.66%   -1.42%     
- Complexity     5126     5293     +167     
============================================
  Files           473      493      +20     
  Lines         21970    23208    +1238     
  Branches       2351     2530     +179     
============================================
+ Hits          15838    16401     +563     
- Misses         4925     5531     +606     
- Partials       1207     1276      +69     

Impacted Files	Coverage Δ
api/src/main/java/ai/djl/modality/cv/Image.java	69.23% <ø> (-4.11%)	⬇️
...i/djl/modality/cv/translator/BigGANTranslator.java	21.42% <ø> (-5.24%)	⬇️
...odality/cv/translator/BigGANTranslatorFactory.java	33.33% <0.00%> (+8.33%)	⬆️
...nslator/InstanceSegmentationTranslatorFactory.java	14.28% <0.00%> (-3.90%)	⬇️
.../modality/cv/translator/YoloTranslatorFactory.java	8.33% <0.00%> (-1.67%)	⬇️
...i/djl/modality/cv/translator/YoloV5Translator.java	5.69% <0.00%> (ø)
...odality/cv/translator/YoloV5TranslatorFactory.java	8.33% <0.00%> (-1.67%)	⬇️
...modality/nlp/embedding/TrainableWordEmbedding.java	45.00% <0.00%> (-3.34%)	⬇️
...pi/src/main/java/ai/djl/ndarray/BytesSupplier.java	54.54% <0.00%> (-12.13%)	⬇️
...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java	15.23% <0.00%> (-0.83%)	⬇️
... and 152 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6c7acba...ad6efd3. Read the comment docs.