test(): added interpolate model for engine serialization testing

Signed-off-by: Abhiram Iyer <[email protected]>

Signed-off-by: Abhiram Iyer <[email protected]>

Author: abhi-iyer

tests/modules/hub.py

@@ -1,4 +1,6 @@
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 import torchvision.models as models
 
 models = {
@@ -65,6 +67,7 @@
     }
 }
 
+# Download sample models
 for n, m in models.items():
     print("Downloading {}".format(n))
     m["model"] = m["model"].eval().cuda()
@@ -74,4 +77,20 @@
         torch.jit.save(trace_model, n + '_traced.jit.pt')
     if m["path"] == "both" or  m["path"] == "script":
         script_model = torch.jit.script(m["model"])
-        torch.jit.save(script_model, n + '_scripted.jit.pt')
+        torch.jit.save(script_model, n + '_scripted.jit.pt')
+
+# Sample Interpolation Model (align_corners=False, for Testing Interpolate Plugin specifically)
+class Interpolate(nn.Module):
+    def __init__(self):
+        super(Interpolate, self).__init__()
+
+    def forward(self, x):
+        return F.interpolate(x, size=(10, 10, 10), align_corners=False, mode="trilinear")
+
+model = Interpolate().eval().cuda()
+x = torch.ones([1, 3, 5, 5, 5]).cuda()
+
+trace_model = torch.jit.trace(model, x)
+torch.jit.save(trace_model, "interpolate_traced.jit.pt")
+
+