fix: bug in vgg16_fp8_ptq example (#2950)

Author: zewenli98

examples/dynamo/vgg16_fp8_ptq.py

@@ -155,7 +155,11 @@ def vgg16(num_classes=1000, init_weights=False):
     ),
 )
 training_dataloader = torch.utils.data.DataLoader(
-    training_dataset, batch_size=args.batch_size, shuffle=True, num_workers=2
+    training_dataset,
+    batch_size=args.batch_size,
+    shuffle=True,
+    num_workers=2,
+    drop_last=True,
 )
 
 data = iter(training_dataloader)
@@ -211,8 +215,12 @@ def calibrate_loop(model):
 )
 
 testing_dataloader = torch.utils.data.DataLoader(
-    testing_dataset, batch_size=args.batch_size, shuffle=False, num_workers=2
-)
+    testing_dataset,
+    batch_size=args.batch_size,
+    shuffle=False,
+    num_workers=2,
+    drop_last=True,
+)  # set drop_last=True to drop the last incomplete batch for static shape `torchtrt.dynamo.compile()`
 
 with torch.no_grad():
     with export_torch_mode():
@@ -235,10 +243,9 @@ def calibrate_loop(model):
         loss = 0.0
         class_probs = []
         class_preds = []
-        model.eval()
         for data, labels in testing_dataloader:
             data, labels = data.cuda(), labels.cuda(non_blocking=True)
-            out = model(data)
+            out = trt_model(data)
             loss += crit(out, labels)
             preds = torch.max(out, 1)[1]
             class_probs.append([F.softmax(i, dim=0) for i in out])