Spaces:

jesseplusplus
/

easy-translate

Running

Iker commited on Apr 26, 2022

Commit

dacb273

1 Parent(s): a911872

Fix bug

Files changed (1) hide show

translate_troch2trt.py CHANGED Viewed

@@ -55,8 +55,10 @@ def main(
         device = "cuda"
         from torch2trt import torch2trt
         model = torch2trt(
-            model.to(device, dtype=dtype),
             [torch.randn((batch_size, max_length)).to(device, dtype=torch.long)],
         )

         device = "cuda"
         from torch2trt import torch2trt
+        model.to(device, dtype=dtype)
         model = torch2trt(
+            model,
             [torch.randn((batch_size, max_length)).to(device, dtype=torch.long)],
         )