📝 Add CUDA examples

francois-rozet · Dec 10, 2020 · a78e4e3 · a78e4e3
1 parent 0e89000
commit a78e4e3
Show file tree

Hide file tree

Showing 6 changed files with 16 additions and 16 deletions.
diff --git a/piqa/gmsd.py b/piqa/gmsd.py
@@ -94,8 +94,8 @@ class GMSD(nn.Module):
 
     Example:
         >>> criterion = GMSD()
-        >>> x = torch.rand(5, 3, 256, 256)
-        >>> y = torch.rand(5, 3, 256, 256)
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
+        >>> y = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x, y)
         >>> l.size()
         torch.Size([])

diff --git a/piqa/lpips.py b/piqa/lpips.py
@@ -88,9 +88,9 @@ class LPIPS(nn.Module):
         use the `torch.no_grad()` context or freeze the weights.
 
     Example:
-        >>> criterion = LPIPS()
-        >>> x = torch.rand(5, 3, 256, 256)
-        >>> y = torch.rand(5, 3, 256, 256)
+        >>> criterion = LPIPS().cuda()
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
+        >>> y = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x, y)
         >>> l.size()
         torch.Size([])

diff --git a/piqa/mdsi.py b/piqa/mdsi.py
@@ -137,8 +137,8 @@ class MDSI(nn.Module):
 
     Example:
         >>> criterion = MDSI()
-        >>> x = torch.rand(5, 3, 256, 256)
-        >>> y = torch.rand(5, 3, 256, 256)
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
+        >>> y = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x, y)
         >>> l.size()
         torch.Size([])

diff --git a/piqa/psnr.py b/piqa/psnr.py
@@ -62,8 +62,8 @@ class PSNR(nn.Module):
 
     Example:
         >>> criterion = PSNR()
-        >>> x = torch.rand(5, 3, 256, 256)
-        >>> y = torch.rand(5, 3, 256, 256)
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
+        >>> y = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x, y)
         >>> l.size()
         torch.Size([])

diff --git a/piqa/ssim.py b/piqa/ssim.py
@@ -226,9 +226,9 @@ class SSIM(nn.Module):
         * Output: (N,) or (1,) depending on `reduction`
 
     Example:
-        >>> criterion = SSIM()
-        >>> x = torch.rand(5, 3, 256, 256)
-        >>> y = torch.rand(5, 3, 256, 256)
+        >>> criterion = SSIM().cuda()
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
+        >>> y = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x, y)
         >>> l.size()
         torch.Size([])
@@ -282,9 +282,9 @@ class MSSSIM(SSIM):
         * Output: (N,) or (1,) depending on `reduction`
 
     Example:
-        >>> criterion = MSSSIM()
-        >>> x = torch.rand(5, 3, 256, 256)
-        >>> y = torch.rand(5, 3, 256, 256)
+        >>> criterion = MSSSIM().cuda()
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
+        >>> y = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x, y)
         >>> l.size()
         torch.Size([])

diff --git a/piqa/tv.py b/piqa/tv.py
@@ -60,7 +60,7 @@ class TV(nn.Module):
 
     Example:
         >>> criterion = TV()
-        >>> x = torch.rand(5, 3, 256, 256)
+        >>> x = torch.rand(5, 3, 256, 256).cuda()
         >>> l = criterion(x)
         >>> l.size()
         torch.Size([])