Update infer_api.py

infer_api.py

@@ -106,7 +106,8 @@ import torch
 from typing import Tuple
 
 @spaces.GPU
-def _warmup(glctx, device=None):
+def _warmup(device=None):
+    glctx = dr.RasterizeCudaContext(device=None)
     device = 'cuda' if device is None else device
     #windows workaround for https://github.com/NVlabs/nvdiffrast/issues/59
     def tensor(*args, **kwargs):
@@ -115,8 +116,7 @@ def _warmup(glctx, device=None):
     tri = tensor([[0, 1, 2]], dtype=torch.int32)
     dr.rasterize(glctx, pos, tri, resolution=[256, 256])
 
-_glctx = dr.RasterizeCudaContext(device=None)
-_warmup(_glctx, device)
+_warmup(device)
 
 #### TEST END ####