Update remove_bg.py

remove_bg.py

@@ -36,9 +36,9 @@ class RemoveBackground(object):
         input_batch = input_tensor.unsqueeze(0)  # create a mini-batch as expected by the model
 
         # move the input and model to GPU for speed if available
-        #if torch.cuda.is_available():
-         #  input_batch = input_batch.to('cuda')
-          # self.model.to('cuda')
+        if torch.cuda.is_available():
+          input_batch = input_batch.to('cuda')
+          self.model.to('cuda')
 
         with torch.no_grad():
             output = self.model(input_batch)['out'][0]