diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bac2866
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,23 @@
+
+images/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# PyTorch weights
+*.tar
+*.pth
+*.gz
+Untitled.ipynb
+Testing notebook.ipynb
diff --git a/image_demo.py b/image_demo.py
index 31e1b92..2dc3bfa 100644
--- a/image_demo.py
+++ b/image_demo.py
@@ -18,7 +18,8 @@
 
 def main():
     model = posenet.load_model(args.model)
-    model = model.cuda()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device=device)
     output_stride = model.output_stride
 
     if args.output_dir:
@@ -34,7 +35,7 @@ def main():
             f, scale_factor=args.scale_factor, output_stride=output_stride)
 
         with torch.no_grad():
-            input_image = torch.Tensor(input_image).cuda()
+            input_image = torch.Tensor(input_image).to(device=device)
 
             heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = model(input_image)
 
diff --git a/webcam_demo.py b/webcam_demo.py
index c2b2139..de77a99 100644
--- a/webcam_demo.py
+++ b/webcam_demo.py
@@ -16,7 +16,8 @@
 
 def main():
     model = posenet.load_model(args.model)
-    model = model.cuda()
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device=device)
     output_stride = model.output_stride
 
     cap = cv2.VideoCapture(args.cam_id)
@@ -30,7 +31,7 @@ def main():
             cap, scale_factor=args.scale_factor, output_stride=output_stride)
 
         with torch.no_grad():
-            input_image = torch.Tensor(input_image).cuda()
+            input_image = torch.Tensor(input_image).to(device=device)
 
             heatmaps_result, offsets_result, displacement_fwd_result, displacement_bwd_result = model(input_image)
 
@@ -59,4 +60,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()