diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal.xcodeproj/project.pbxproj b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal.xcodeproj/project.pbxproj
index fb51907..1228aff 100644
--- a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal.xcodeproj/project.pbxproj
+++ b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal.xcodeproj/project.pbxproj
@@ -7,19 +7,33 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		492EFF2688481D0773EA77DA /* libPods-HelloWorld-Metal.a in Frameworks */ = {isa = PBXBuildFile; fileRef = C8FA3F26EDC2BFF1CFF70E63 /* libPods-HelloWorld-Metal.a */; };
 		8BA1BEFC26C1C13C008AD5B2 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8BA1BEFB26C1C13C008AD5B2 /* AppDelegate.swift */; };
 		8BA1BF0026C1C13C008AD5B2 /* ViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8BA1BEFF26C1C13C008AD5B2 /* ViewController.swift */; };
 		8BA1BF0326C1C13C008AD5B2 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF0126C1C13C008AD5B2 /* Main.storyboard */; };
 		8BA1BF0526C1C13D008AD5B2 /* Assets.xcassets in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF0426C1C13D008AD5B2 /* Assets.xcassets */; };
 		8BA1BF0826C1C13D008AD5B2 /* LaunchScreen.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF0626C1C13D008AD5B2 /* LaunchScreen.storyboard */; };
-		8BA1BF1026C1C2CC008AD5B2 /* image.png in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF0F26C1C2CC008AD5B2 /* image.png */; };
+		8BA1BF1026C1C2CC008AD5B2 /* wolf.png in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF0F26C1C2CC008AD5B2 /* wolf.png */; };
 		8BA1BF1726C1C305008AD5B2 /* words.txt in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF1526C1C305008AD5B2 /* words.txt */; };
 		8BA1BF1E26C1C356008AD5B2 /* UIImage+Helper.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8BA1BF1D26C1C356008AD5B2 /* UIImage+Helper.swift */; };
 		8BA1BF2926C1C885008AD5B2 /* TorchModule.mm in Sources */ = {isa = PBXBuildFile; fileRef = 8BA1BF2726C1C885008AD5B2 /* TorchModule.mm */; };
 		8BA1BF4B26C1DDCD008AD5B2 /* model.pt in Resources */ = {isa = PBXBuildFile; fileRef = 8BA1BF4A26C1DDCD008AD5B2 /* model.pt */; };
+		BDCC15C627578D00005A979D /* animal-3.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15BA27578CFF005A979D /* animal-3.jpg */; };
+		BDCC15C727578D00005A979D /* food-1.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15BB27578CFF005A979D /* food-1.jpg */; };
+		BDCC15C827578D00005A979D /* car-3.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15BC27578CFF005A979D /* car-3.jpg */; };
+		BDCC15C927578D00005A979D /* girl-1.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15BD27578CFF005A979D /* girl-1.jpg */; };
+		BDCC15CA27578D00005A979D /* girl-2.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15BE27578CFF005A979D /* girl-2.jpg */; };
+		BDCC15CC27578D00005A979D /* car-1.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15C027578CFF005A979D /* car-1.jpg */; };
+		BDCC15CD27578D00005A979D /* car-2.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15C127578D00005A979D /* car-2.jpg */; };
+		BDCC15CE27578D00005A979D /* animal-2.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15C227578D00005A979D /* animal-2.jpg */; };
+		BDCC15CF27578D00005A979D /* kitten_small.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15C327578D00005A979D /* kitten_small.jpg */; };
+		BDCC15D027578D00005A979D /* girl-3.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15C427578D00005A979D /* girl-3.jpg */; };
+		BDCC15D127578D00005A979D /* animal-1.jpg in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15C527578D00005A979D /* animal-1.jpg */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXFileReference section */
+		2F88BD040F1A45FD07B00711 /* Pods-HelloWorld-Metal.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-HelloWorld-Metal.release.xcconfig"; path = "Target Support Files/Pods-HelloWorld-Metal/Pods-HelloWorld-Metal.release.xcconfig"; sourceTree = "<group>"; };
+		5D06AFB6753649843FB01282 /* Pods-HelloWorld-Metal.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-HelloWorld-Metal.debug.xcconfig"; path = "Target Support Files/Pods-HelloWorld-Metal/Pods-HelloWorld-Metal.debug.xcconfig"; sourceTree = "<group>"; };
 		8BA1BEF826C1C13C008AD5B2 /* HelloWorld-Metal.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = "HelloWorld-Metal.app"; sourceTree = BUILT_PRODUCTS_DIR; };
 		8BA1BEFB26C1C13C008AD5B2 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = "<group>"; };
 		8BA1BEFF26C1C13C008AD5B2 /* ViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ViewController.swift; sourceTree = "<group>"; };
@@ -27,13 +41,25 @@
 		8BA1BF0426C1C13D008AD5B2 /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
 		8BA1BF0726C1C13D008AD5B2 /* Base */ = {isa = PBXFileReference; lastKnownFileType = file.storyboard; name = Base; path = Base.lproj/LaunchScreen.storyboard; sourceTree = "<group>"; };
 		8BA1BF0926C1C13D008AD5B2 /* Info.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = Info.plist; sourceTree = "<group>"; };
-		8BA1BF0F26C1C2CC008AD5B2 /* image.png */ = {isa = PBXFileReference; lastKnownFileType = image.png; path = image.png; sourceTree = "<group>"; };
+		8BA1BF0F26C1C2CC008AD5B2 /* wolf.png */ = {isa = PBXFileReference; lastKnownFileType = image.png; path = wolf.png; sourceTree = "<group>"; };
 		8BA1BF1526C1C305008AD5B2 /* words.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = words.txt; sourceTree = "<group>"; };
 		8BA1BF1D26C1C356008AD5B2 /* UIImage+Helper.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "UIImage+Helper.swift"; sourceTree = "<group>"; };
 		8BA1BF2126C1C861008AD5B2 /* HelloWorld-Metal-Bridging-Header.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = "HelloWorld-Metal-Bridging-Header.h"; sourceTree = "<group>"; };
 		8BA1BF2726C1C885008AD5B2 /* TorchModule.mm */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.cpp.objcpp; path = TorchModule.mm; sourceTree = "<group>"; };
 		8BA1BF2826C1C885008AD5B2 /* TorchModule.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = TorchModule.h; sourceTree = "<group>"; };
 		8BA1BF4A26C1DDCD008AD5B2 /* model.pt */ = {isa = PBXFileReference; lastKnownFileType = file; path = model.pt; sourceTree = "<group>"; };
+		BDCC15BA27578CFF005A979D /* animal-3.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "animal-3.jpg"; sourceTree = "<group>"; };
+		BDCC15BB27578CFF005A979D /* food-1.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "food-1.jpg"; sourceTree = "<group>"; };
+		BDCC15BC27578CFF005A979D /* car-3.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "car-3.jpg"; sourceTree = "<group>"; };
+		BDCC15BD27578CFF005A979D /* girl-1.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "girl-1.jpg"; sourceTree = "<group>"; };
+		BDCC15BE27578CFF005A979D /* girl-2.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "girl-2.jpg"; sourceTree = "<group>"; };
+		BDCC15C027578CFF005A979D /* car-1.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "car-1.jpg"; sourceTree = "<group>"; };
+		BDCC15C127578D00005A979D /* car-2.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "car-2.jpg"; sourceTree = "<group>"; };
+		BDCC15C227578D00005A979D /* animal-2.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "animal-2.jpg"; sourceTree = "<group>"; };
+		BDCC15C327578D00005A979D /* kitten_small.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = kitten_small.jpg; sourceTree = "<group>"; };
+		BDCC15C427578D00005A979D /* girl-3.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "girl-3.jpg"; sourceTree = "<group>"; };
+		BDCC15C527578D00005A979D /* animal-1.jpg */ = {isa = PBXFileReference; lastKnownFileType = image.jpeg; path = "animal-1.jpg"; sourceTree = "<group>"; };
+		C8FA3F26EDC2BFF1CFF70E63 /* libPods-HelloWorld-Metal.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-HelloWorld-Metal.a"; sourceTree = BUILT_PRODUCTS_DIR; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -41,17 +67,29 @@
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				492EFF2688481D0773EA77DA /* libPods-HelloWorld-Metal.a in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
 /* End PBXFrameworksBuildPhase section */
 
 /* Begin PBXGroup section */
+		6934F693EDBC7C05977F1E39 /* Pods */ = {
+			isa = PBXGroup;
+			children = (
+				5D06AFB6753649843FB01282 /* Pods-HelloWorld-Metal.debug.xcconfig */,
+				2F88BD040F1A45FD07B00711 /* Pods-HelloWorld-Metal.release.xcconfig */,
+			);
+			path = Pods;
+			sourceTree = "<group>";
+		};
 		8BA1BEEF26C1C13C008AD5B2 = {
 			isa = PBXGroup;
 			children = (
 				8BA1BEFA26C1C13C008AD5B2 /* HelloWorld-Metal */,
 				8BA1BEF926C1C13C008AD5B2 /* Products */,
+				6934F693EDBC7C05977F1E39 /* Pods */,
+				D598EFAF0D884557076C0509 /* Frameworks */,
 			);
 			sourceTree = "<group>";
 		};
@@ -66,7 +104,7 @@
 		8BA1BEFA26C1C13C008AD5B2 /* HelloWorld-Metal */ = {
 			isa = PBXGroup;
 			children = (
-				8BA1BF0F26C1C2CC008AD5B2 /* image.png */,
+				BDCC15B927578C5D005A979D /* images */,
 				8BA1BF1326C1C305008AD5B2 /* model */,
 				8BA1BF2526C1C885008AD5B2 /* TorchBridge */,
 				8BA1BEFB26C1C13C008AD5B2 /* AppDelegate.swift */,
@@ -99,6 +137,33 @@
 			path = TorchBridge;
 			sourceTree = "<group>";
 		};
+		BDCC15B927578C5D005A979D /* images */ = {
+			isa = PBXGroup;
+			children = (
+				8BA1BF0F26C1C2CC008AD5B2 /* wolf.png */,
+				BDCC15C527578D00005A979D /* animal-1.jpg */,
+				BDCC15C227578D00005A979D /* animal-2.jpg */,
+				BDCC15BA27578CFF005A979D /* animal-3.jpg */,
+				BDCC15C027578CFF005A979D /* car-1.jpg */,
+				BDCC15C127578D00005A979D /* car-2.jpg */,
+				BDCC15BC27578CFF005A979D /* car-3.jpg */,
+				BDCC15BB27578CFF005A979D /* food-1.jpg */,
+				BDCC15BD27578CFF005A979D /* girl-1.jpg */,
+				BDCC15BE27578CFF005A979D /* girl-2.jpg */,
+				BDCC15C427578D00005A979D /* girl-3.jpg */,
+				BDCC15C327578D00005A979D /* kitten_small.jpg */,
+			);
+			path = images;
+			sourceTree = "<group>";
+		};
+		D598EFAF0D884557076C0509 /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				C8FA3F26EDC2BFF1CFF70E63 /* libPods-HelloWorld-Metal.a */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
 /* End PBXGroup section */
 
 /* Begin PBXNativeTarget section */
@@ -106,6 +171,7 @@
 			isa = PBXNativeTarget;
 			buildConfigurationList = 8BA1BF0C26C1C13D008AD5B2 /* Build configuration list for PBXNativeTarget "HelloWorld-Metal" */;
 			buildPhases = (
+				DD3150C9787B8E05005715ED /* [CP] Check Pods Manifest.lock */,
 				8BA1BEF426C1C13C008AD5B2 /* Sources */,
 				8BA1BEF526C1C13C008AD5B2 /* Frameworks */,
 				8BA1BEF626C1C13C008AD5B2 /* Resources */,
@@ -159,16 +225,52 @@
 			buildActionMask = 2147483647;
 			files = (
 				8BA1BF0826C1C13D008AD5B2 /* LaunchScreen.storyboard in Resources */,
+				BDCC15C927578D00005A979D /* girl-1.jpg in Resources */,
 				8BA1BF0526C1C13D008AD5B2 /* Assets.xcassets in Resources */,
+				BDCC15CA27578D00005A979D /* girl-2.jpg in Resources */,
+				BDCC15CC27578D00005A979D /* car-1.jpg in Resources */,
+				BDCC15D127578D00005A979D /* animal-1.jpg in Resources */,
 				8BA1BF0326C1C13C008AD5B2 /* Main.storyboard in Resources */,
+				BDCC15CF27578D00005A979D /* kitten_small.jpg in Resources */,
+				BDCC15C627578D00005A979D /* animal-3.jpg in Resources */,
+				BDCC15D027578D00005A979D /* girl-3.jpg in Resources */,
 				8BA1BF1726C1C305008AD5B2 /* words.txt in Resources */,
-				8BA1BF1026C1C2CC008AD5B2 /* image.png in Resources */,
+				BDCC15CE27578D00005A979D /* animal-2.jpg in Resources */,
+				8BA1BF1026C1C2CC008AD5B2 /* wolf.png in Resources */,
 				8BA1BF4B26C1DDCD008AD5B2 /* model.pt in Resources */,
+				BDCC15CD27578D00005A979D /* car-2.jpg in Resources */,
+				BDCC15C827578D00005A979D /* car-3.jpg in Resources */,
+				BDCC15C727578D00005A979D /* food-1.jpg in Resources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
 /* End PBXResourcesBuildPhase section */
 
+/* Begin PBXShellScriptBuildPhase section */
+		DD3150C9787B8E05005715ED /* [CP] Check Pods Manifest.lock */ = {
+			isa = PBXShellScriptBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			inputFileListPaths = (
+			);
+			inputPaths = (
+				"${PODS_PODFILE_DIR_PATH}/Podfile.lock",
+				"${PODS_ROOT}/Manifest.lock",
+			);
+			name = "[CP] Check Pods Manifest.lock";
+			outputFileListPaths = (
+			);
+			outputPaths = (
+				"$(DERIVED_FILE_DIR)/Pods-HelloWorld-Metal-checkManifestLockResult.txt",
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+			shellPath = /bin/sh;
+			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
+			showEnvVarsInLog = 0;
+		};
+/* End PBXShellScriptBuildPhase section */
+
 /* Begin PBXSourcesBuildPhase section */
 		8BA1BEF426C1C13C008AD5B2 /* Sources */ = {
 			isa = PBXSourcesBuildPhase;
@@ -321,15 +423,17 @@
 		};
 		8BA1BF0D26C1C13D008AD5B2 /* Debug */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = 5D06AFB6753649843FB01282 /* Pods-HelloWorld-Metal.debug.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_IDENTITY = "Apple Development";
-				CODE_SIGN_STYLE = Manual;
-				DEVELOPMENT_TEAM = "";
+				CODE_SIGN_STYLE = Automatic;
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = "HelloWorld-Metal/Info.plist";
+				IPHONEOS_DEPLOYMENT_TARGET = 14.0;
 				LD_RUNPATH_SEARCH_PATHS = (
 					"$(inherited)",
 					"@executable_path/Frameworks",
@@ -346,15 +450,17 @@
 		};
 		8BA1BF0E26C1C13D008AD5B2 /* Release */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = 2F88BD040F1A45FD07B00711 /* Pods-HelloWorld-Metal.release.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_IDENTITY = "Apple Development";
-				CODE_SIGN_STYLE = Manual;
-				DEVELOPMENT_TEAM = "";
+				CODE_SIGN_STYLE = Automatic;
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = "HelloWorld-Metal/Info.plist";
+				IPHONEOS_DEPLOYMENT_TARGET = 14.0;
 				LD_RUNPATH_SEARCH_PATHS = (
 					"$(inherited)",
 					"@executable_path/Frameworks",
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/Base.lproj/Main.storyboard b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/Base.lproj/Main.storyboard
index 52e0d13..d42fc05 100644
--- a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/Base.lproj/Main.storyboard
+++ b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/Base.lproj/Main.storyboard
@@ -1,9 +1,9 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="18122" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="MJX-KQ-XfU">
+<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="19455" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="MJX-KQ-XfU">
     <device id="retina6_1" orientation="portrait" appearance="light"/>
     <dependencies>
         <deployment identifier="iOS"/>
-        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="18093"/>
+        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="19454"/>
         <capability name="Safe area layout guides" minToolsVersion="9.0"/>
         <capability name="System colors in document resources" minToolsVersion="11.0"/>
         <capability name="documents saved in the Xcode 8 format" minToolsVersion="8.0"/>
@@ -28,7 +28,7 @@
         <!--Image Classification-->
         <scene sceneID="tne-QT-ifu">
             <objects>
-                <viewController id="BYZ-38-t0r" customClass="ViewController" customModule="HelloWorldMetal" customModuleProvider="target" sceneMemberID="viewController">
+                <viewController id="BYZ-38-t0r" customClass="ViewController" customModule="HelloWorld_Metal" customModuleProvider="target" sceneMemberID="viewController">
                     <view key="view" contentMode="scaleToFill" id="8bC-Xf-vdC">
                         <rect key="frame" x="0.0" y="0.0" width="414" height="896"/>
                         <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
@@ -39,40 +39,92 @@
                                     <constraint firstAttribute="width" secondItem="tcN-Mf-gVj" secondAttribute="height" multiplier="1:1" id="9wt-9w-9ry"/>
                                 </constraints>
                             </imageView>
-                            <textView clipsSubviews="YES" multipleTouchEnabled="YES" contentMode="scaleToFill" editable="NO" textAlignment="natural" selectable="NO" translatesAutoresizingMaskIntoConstraints="NO" id="E7l-Rs-UJz">
-                                <rect key="frame" x="20" y="502" width="374" height="340"/>
+                            <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="aAV-kg-ted">
+                                <rect key="frame" x="172" y="831" width="70.5" height="31"/>
+                                <state key="normal" title="Button"/>
+                                <buttonConfiguration key="configuration" style="plain" title="Predict"/>
+                                <connections>
+                                    <action selector="runPredict:" destination="BYZ-38-t0r" eventType="touchUpInside" id="pXp-cH-wQb"/>
+                                </connections>
+                            </button>
+                            <stepper opaque="NO" contentMode="scaleToFill" horizontalHuggingPriority="750" verticalHuggingPriority="750" contentHorizontalAlignment="center" contentVerticalAlignment="center" autorepeat="NO" wraps="YES" maximumValue="100" translatesAutoresizingMaskIntoConstraints="NO" id="qy4-9O-Fqi">
+                                <rect key="frame" x="160" y="490" width="94" height="32"/>
+                                <connections>
+                                    <action selector="changeImage:" destination="BYZ-38-t0r" eventType="valueChanged" id="QnP-Mj-Zyu"/>
+                                </connections>
+                            </stepper>
+                            <tableView clipsSubviews="YES" contentMode="scaleToFill" alwaysBounceVertical="YES" dataMode="prototypes" style="plain" separatorStyle="none" allowsSelection="NO" rowHeight="45" estimatedRowHeight="-1" sectionHeaderHeight="-1" estimatedSectionHeaderHeight="-1" sectionFooterHeight="-1" estimatedSectionFooterHeight="-1" translatesAutoresizingMaskIntoConstraints="NO" id="Pdd-lq-6jU">
+                                <rect key="frame" x="20" y="559" width="374" height="264"/>
                                 <color key="backgroundColor" systemColor="systemBackgroundColor"/>
-                                <color key="textColor" systemColor="labelColor"/>
-                                <fontDescription key="fontDescription" type="system" pointSize="18"/>
-                                <textInputTraits key="textInputTraits" autocapitalizationType="sentences"/>
-                            </textView>
+                                <prototypes>
+                                    <tableViewCell clipsSubviews="YES" contentMode="scaleToFill" insetsLayoutMarginsFromSafeArea="NO" selectionStyle="blue" hidesAccessoryWhenEditing="NO" indentationLevel="1" indentationWidth="0.0" reuseIdentifier="result" textLabel="PCe-zs-Bvh" detailTextLabel="489-Jh-vGK" rowHeight="46" style="IBUITableViewCellStyleSubtitle" id="Igw-8i-eEj">
+                                        <rect key="frame" x="0.0" y="44.5" width="374" height="46"/>
+                                        <autoresizingMask key="autoresizingMask"/>
+                                        <tableViewCellContentView key="contentView" opaque="NO" clipsSubviews="YES" multipleTouchEnabled="YES" contentMode="center" insetsLayoutMarginsFromSafeArea="NO" tableViewCell="Igw-8i-eEj" id="0a0-zl-Zyy">
+                                            <rect key="frame" x="0.0" y="0.0" width="374" height="46"/>
+                                            <autoresizingMask key="autoresizingMask"/>
+                                            <subviews>
+                                                <label opaque="NO" multipleTouchEnabled="YES" contentMode="left" insetsLayoutMarginsFromSafeArea="NO" text="Title" textAlignment="natural" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" id="PCe-zs-Bvh">
+                                                    <rect key="frame" x="20" y="6.5" width="28.5" height="17"/>
+                                                    <autoresizingMask key="autoresizingMask"/>
+                                                    <fontDescription key="fontDescription" style="UICTFontTextStyleBody"/>
+                                                    <nil key="textColor"/>
+                                                    <nil key="highlightedColor"/>
+                                                </label>
+                                                <label opaque="NO" multipleTouchEnabled="YES" contentMode="left" insetsLayoutMarginsFromSafeArea="NO" text="Subtitle" textAlignment="natural" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" id="489-Jh-vGK">
+                                                    <rect key="frame" x="20" y="25" width="44" height="14.5"/>
+                                                    <autoresizingMask key="autoresizingMask"/>
+                                                    <fontDescription key="fontDescription" style="UICTFontTextStyleFootnote"/>
+                                                    <color key="textColor" systemColor="secondaryLabelColor"/>
+                                                    <nil key="highlightedColor"/>
+                                                </label>
+                                            </subviews>
+                                        </tableViewCellContentView>
+                                    </tableViewCell>
+                                </prototypes>
+                            </tableView>
+                            <label opaque="NO" userInteractionEnabled="NO" contentMode="left" horizontalHuggingPriority="251" verticalHuggingPriority="251" text="Label" textAlignment="natural" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" translatesAutoresizingMaskIntoConstraints="NO" id="gbA-Ug-toR">
+                                <rect key="frame" x="186" y="530" width="42" height="21"/>
+                                <fontDescription key="fontDescription" type="system" pointSize="17"/>
+                                <nil key="textColor"/>
+                                <nil key="highlightedColor"/>
+                            </label>
                         </subviews>
                         <viewLayoutGuide key="safeArea" id="6Tk-OE-BBY"/>
                         <color key="backgroundColor" systemColor="systemBackgroundColor"/>
                         <constraints>
+                            <constraint firstItem="Pdd-lq-6jU" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="3rY-Q5-vyN"/>
+                            <constraint firstItem="aAV-kg-ted" firstAttribute="bottom" secondItem="6Tk-OE-BBY" secondAttribute="bottom" id="4Yg-He-1ZE"/>
                             <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="tcN-Mf-gVj" secondAttribute="trailing" constant="20" id="6qF-bk-R5a"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="E7l-Rs-UJz" secondAttribute="trailing" constant="20" id="7SV-x7-w0l"/>
+                            <constraint firstItem="gbA-Ug-toR" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="8ae-F0-4Zh"/>
+                            <constraint firstItem="qy4-9O-Fqi" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="9wT-68-2w9"/>
                             <constraint firstItem="tcN-Mf-gVj" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="Gig-Ia-Utu"/>
-                            <constraint firstItem="E7l-Rs-UJz" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="JGO-EO-qpf"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="bottom" secondItem="E7l-Rs-UJz" secondAttribute="bottom" constant="20" id="dr2-Xm-IG4"/>
+                            <constraint firstItem="aAV-kg-ted" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="Pb2-2G-usQ"/>
+                            <constraint firstItem="gbA-Ug-toR" firstAttribute="top" secondItem="qy4-9O-Fqi" secondAttribute="bottom" constant="8" symbolic="YES" id="eya-SQ-EHq"/>
+                            <constraint firstItem="aAV-kg-ted" firstAttribute="top" secondItem="Pdd-lq-6jU" secondAttribute="bottom" constant="8" symbolic="YES" id="lBU-d2-eRA"/>
+                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="Pdd-lq-6jU" secondAttribute="trailing" constant="20" id="lpI-D5-kgi"/>
+                            <constraint firstItem="qy4-9O-Fqi" firstAttribute="top" secondItem="tcN-Mf-gVj" secondAttribute="bottom" constant="8" symbolic="YES" id="nZz-51-QLt"/>
                             <constraint firstItem="tcN-Mf-gVj" firstAttribute="top" secondItem="6Tk-OE-BBY" secondAttribute="top" constant="20" id="uBS-lr-amS"/>
-                            <constraint firstItem="E7l-Rs-UJz" firstAttribute="top" secondItem="tcN-Mf-gVj" secondAttribute="bottom" constant="20" id="zVo-eH-khM"/>
+                            <constraint firstItem="Pdd-lq-6jU" firstAttribute="top" secondItem="gbA-Ug-toR" secondAttribute="bottom" constant="8" symbolic="YES" id="xJv-jJ-eEP"/>
                         </constraints>
                     </view>
                     <navigationItem key="navigationItem" title="Image Classification" id="n4K-Na-VKn"/>
                     <connections>
+                        <outlet property="durationLabel" destination="gbA-Ug-toR" id="zvH-Xo-NdQ"/>
+                        <outlet property="imageStepper" destination="qy4-9O-Fqi" id="CRJ-Sz-qW8"/>
                         <outlet property="imageView" destination="tcN-Mf-gVj" id="9y4-1U-aPc"/>
-                        <outlet property="resultView" destination="E7l-Rs-UJz" id="NIv-cW-amD"/>
+                        <outlet property="predictButton" destination="aAV-kg-ted" id="beT-pE-GIZ"/>
+                        <outlet property="resultsView" destination="Pdd-lq-6jU" id="wDn-Fy-cI8"/>
                     </connections>
                 </viewController>
                 <placeholder placeholderIdentifier="IBFirstResponder" id="dkx-z0-nzr" sceneMemberID="firstResponder"/>
             </objects>
-            <point key="canvasLocation" x="-278" y="-86"/>
+            <point key="canvasLocation" x="-278.26086956521743" y="-86.383928571428569"/>
         </scene>
     </scenes>
     <resources>
-        <systemColor name="labelColor">
-            <color white="0.0" alpha="1" colorSpace="custom" customColorSpace="genericGamma22GrayColorSpace"/>
+        <systemColor name="secondaryLabelColor">
+            <color red="0.23529411764705882" green="0.23529411764705882" blue="0.2627450980392157" alpha="0.59999999999999998" colorSpace="custom" customColorSpace="sRGB"/>
         </systemColor>
         <systemColor name="systemBackgroundColor">
             <color white="1" alpha="1" colorSpace="custom" customColorSpace="genericGamma22GrayColorSpace"/>
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.h b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.h
index 883120f..c19c891 100644
--- a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.h
+++ b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.h
@@ -4,11 +4,11 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface TorchModule : NSObject
 
-- (nullable instancetype)initWithFileAtPath:(NSString*)filePath
-    NS_SWIFT_NAME(init(fileAtPath:))NS_DESIGNATED_INITIALIZER;
+- (nullable instancetype)initWithFileAtPath:(NSString*)filePath width:(long)width height:(long)height
+    NS_SWIFT_NAME(init(fileAtPath:width:height:))NS_DESIGNATED_INITIALIZER;
 + (instancetype)new NS_UNAVAILABLE;
 - (instancetype)init NS_UNAVAILABLE;
-- (nullable NSArray<NSNumber*>*)predictImage:(void*)imageBuffer NS_SWIFT_NAME(predict(image:));
+- (nullable float*)predictImage:(void*)imageBuffer NS_SWIFT_NAME(predict(image:));
 
 @end
 
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.mm b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.mm
index 81ef769..5e86f11 100644
--- a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.mm
+++ b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/TorchBridge/TorchModule.mm
@@ -3,14 +3,20 @@
 
 @implementation TorchModule {
  @protected
-  torch::jit::mobile::Module _impl;
+  torch::jit::mobile::Module _model;
+  at::Tensor _output;
+  int64_t _inputSize[4];
 }
 
-- (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
+- (nullable instancetype)initWithFileAtPath:(NSString*)filePath width:(long)width height:(long)height {
   self = [super init];
   if (self) {
+    _inputSize[0] = 1;
+    _inputSize[1] = 3;
+    _inputSize[2] = width;
+    _inputSize[3] = height;
     try {
-      _impl = torch::jit::_load_for_mobile(filePath.UTF8String);
+      _model = torch::jit::_load_for_mobile(filePath.UTF8String);
     } catch (const std::exception& exception) {
       NSLog(@"%s", exception.what());
       return nil;
@@ -19,20 +25,12 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
   return self;
 }
 
-- (NSArray<NSNumber*>*)predictImage:(void*)imageBuffer {
+- (float*)predictImage:(void*)imageBuffer {
   try {
-    c10::InferenceMode mode;
-    at::Tensor tensor = torch::from_blob(imageBuffer, {1, 3, 224, 224}, at::kFloat).metal();
-    auto outputTensor = _impl.forward({tensor}).toTensor().cpu();
-    float* floatBuffer = outputTensor.data_ptr<float>();
-    if (!floatBuffer) {
-      return nil;
-    }
-    NSMutableArray* results = [[NSMutableArray alloc] init];
-    for (int i = 0; i < 1000; i++) {
-      [results addObject:@(floatBuffer[i])];
-    }
-    return [results copy];
+    c10::InferenceMode guard(true);
+    at::Tensor tensor = torch::from_blob(imageBuffer, _inputSize, at::kFloat).metal();
+    _output = _model.forward({tensor}).toTensor().cpu();
+    return _output.data_ptr<float>();
   } catch (const std::exception& exception) {
     NSLog(@"%s", exception.what());
   }
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/UIImage+Helper.swift b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/UIImage+Helper.swift
index d8bab86..c4dbd86 100755
--- a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/UIImage+Helper.swift
+++ b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/UIImage+Helper.swift
@@ -1,47 +1,44 @@
 import UIKit
 
 extension UIImage {
-    func resized(to newSize: CGSize, scale: CGFloat = 1) -> UIImage {
-        let format = UIGraphicsImageRendererFormat.default()
-        format.scale = scale
-        let renderer = UIGraphicsImageRenderer(size: newSize, format: format)
-        let image = renderer.image { _ in
-            draw(in: CGRect(origin: .zero, size: newSize))
-        }
-        return image
+  func resized(to newSize: CGSize, scale: CGFloat = 1) -> UIImage {
+    let format = UIGraphicsImageRendererFormat.default()
+    format.scale = scale
+    let renderer = UIGraphicsImageRenderer(size: newSize, format: format)
+    return renderer.image { _ in draw(in: CGRect(origin: .zero, size: newSize)) }
+  }
+
+  func normalized() -> [Float32]? {
+    guard let cgImage = self.cgImage else { return nil }
+    let w = cgImage.width
+    let h = cgImage.height
+    let bytesPerPixel = 4
+    let bytesPerRow = bytesPerPixel * w
+    let bitsPerComponent = 8
+
+    // Render image into our buffer so we have access to pixel bytes
+    var rawBytes: [UInt8] = [UInt8](repeating: 0, count: w * h * 4)
+    rawBytes.withUnsafeMutableBytes { ptr in
+      if let context = CGContext(data: ptr.baseAddress,
+                                 width: w,
+                                 height: h,
+                                 bitsPerComponent: bitsPerComponent,
+                                 bytesPerRow: bytesPerRow,
+                                 space: CGColorSpaceCreateDeviceRGB(),
+                                 bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) {
+        let rect = CGRect(x: 0, y: 0, width: w, height: h)
+        context.draw(cgImage, in: rect)
+      }
     }
 
-    func normalized() -> [Float32]? {
-        guard let cgImage = self.cgImage else {
-            return nil
-        }
-        let w = cgImage.width
-        let h = cgImage.height
-        let bytesPerPixel = 4
-        let bytesPerRow = bytesPerPixel * w
-        let bitsPerComponent = 8
-        var rawBytes: [UInt8] = [UInt8](repeating: 0, count: w * h * 4)
-        rawBytes.withUnsafeMutableBytes { ptr in
-            if let cgImage = self.cgImage,
-                let context = CGContext(data: ptr.baseAddress,
-                                        width: w,
-                                        height: h,
-                                        bitsPerComponent: bitsPerComponent,
-                                        bytesPerRow: bytesPerRow,
-                                        space: CGColorSpaceCreateDeviceRGB(),
-                                        bitmapInfo: CGImageAlphaInfo.premultipliedLast.rawValue) {
-                let rect = CGRect(x: 0, y: 0, width: w, height: h)
-                context.draw(cgImage, in: rect)
-            }
-        }
-        var normalizedBuffer: [Float32] = [Float32](repeating: 0, count: w * h * 3)
-        // normalize the pixel buffer
-        // see https://pytorch.org/hub/pytorch_vision_resnet/ for more detail
-        for i in 0 ..< w * h {
-            normalizedBuffer[i] = (Float32(rawBytes[i * 4 + 0]) / 255.0 - 0.485) / 0.229 // R
-            normalizedBuffer[w * h + i] = (Float32(rawBytes[i * 4 + 1]) / 255.0 - 0.456) / 0.224 // G
-            normalizedBuffer[w * h * 2 + i] = (Float32(rawBytes[i * 4 + 2]) / 255.0 - 0.406) / 0.225 // B
-        }
-        return normalizedBuffer
+    var normalizedBuffer: [Float32] = [Float32](repeating: 0, count: w * h * 3)
+    // normalize the pixel buffer
+    // see https://pytorch.org/hub/pytorch_vision_resnet/ for more detail
+    for index in 0 ..< w * h {
+      normalizedBuffer[index] = (Float32(rawBytes[index * 4 + 0]) / 255.0 - 0.485) / 0.229 // R
+      normalizedBuffer[w * h + index] = (Float32(rawBytes[index * 4 + 1]) / 255.0 - 0.456) / 0.224 // G
+      normalizedBuffer[w * h * 2 + index] = (Float32(rawBytes[index * 4 + 2]) / 255.0 - 0.406) / 0.225 // B
     }
+    return normalizedBuffer
+  }
 }
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift
index 96894eb..d4b8fff 100644
--- a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift
+++ b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/ViewController.swift
@@ -1,43 +1,124 @@
 import UIKit
 
-class ViewController: UIViewController {
-    @IBOutlet var imageView: UIImageView!
-    @IBOutlet var resultView: UITextView!
-    private lazy var module: TorchModule = {
-        if let filePath = Bundle.main.path(forResource: "model", ofType: "pt"),
-            let module = TorchModule(fileAtPath: filePath) {
-            return module
-        } else {
-            fatalError("Can't find the model file!")
-        }
-    }()
-
-    private lazy var labels: [String] = {
-        if let filePath = Bundle.main.path(forResource: "words", ofType: "txt"),
-            let labels = try? String(contentsOfFile: filePath) {
-            return labels.components(separatedBy: .newlines)
-        } else {
-            fatalError("Can't find the text file!")
-        }
-    }()
-
-    override func viewDidLoad() {
-        super.viewDidLoad()
-        let image = UIImage(named: "image.png")!
-        imageView.image = image
-        let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
-        guard var pixelBuffer = resizedImage.normalized() else {
-            return
-        }
-        guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
-            return
-        }
-        let zippedResults = zip(labels.indices, outputs)
-        let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
-        var text = ""
-        for result in sortedResults {
-            text += "\u{2022} \(labels[result.0]) \n\n"
+struct Prediction {
+  let confidence: Float
+  let label: String
+}
+
+class ViewController: UIViewController, UITableViewDataSource, UITableViewDelegate {
+  @IBOutlet weak var imageView: UIImageView!
+  @IBOutlet weak var predictButton: UIButton!
+  @IBOutlet weak var imageStepper: UIStepper!
+  @IBOutlet weak var durationLabel: UILabel!
+  @IBOutlet weak var resultsView: UITableView!
+
+  var predictions = [Prediction]()
+
+  let width: Int = 224
+  let height: Int = 224
+
+  private lazy var module: TorchModule = {
+    guard let filePath = Bundle.main.path(forResource: "model", ofType: "pt"),
+          let module = TorchModule(fileAtPath: filePath, width: width, height: height)
+    else {
+      fatalError("Can't find the model file!")
+    }
+    return module
+  }()
+
+  private lazy var labels: [Substring] = {
+    guard let filePath = Bundle.main.path(forResource: "words", ofType: "txt"),
+          let labels = try? String(contentsOfFile: filePath)
+    else {
+      fatalError("Can't find the text file!")
+    }
+    return labels.split(whereSeparator: \.isNewline)
+  }()
+
+  private lazy var images: [UIImage?] = {
+    let imagePaths = Bundle.main.paths(forResourcesOfType: "jpg", inDirectory: nil) +
+    Bundle.main.paths(forResourcesOfType: "png", inDirectory: nil)
+    print(imagePaths)
+    return imagePaths.map { UIImage(named: $0) }
+  }()
+
+  private var activeImage: UIImage? { images[Int(imageStepper.value)] }
+
+  override func viewDidLoad() {
+    super.viewDidLoad()
+    resultsView.dataSource = self
+    imageStepper.value = 0
+    imageStepper.minimumValue = 0
+    imageStepper.maximumValue = Double(images.count - 1)
+    updateImageView()
+    doPrediction()
+  }
+
+  func updateImageView() {
+    imageView.image = images[Int(imageStepper.value)]
+  }
+
+  private var imageFloats: [Float32]? {
+    guard let resizedImage = activeImage?.resized(to: CGSize(width: width, height: height)) else { return nil }
+    return resizedImage.normalized()
+  }
+
+  func doPrediction() {
+    guard var input = self.imageFloats else { return }
+    predictions.removeAll()
+    resultsView.reloadData()
+
+    input.withUnsafeMutableBufferPointer { ptr in
+      let startTime: DispatchTime = .now()
+      if let output = module.predict(image: &(ptr.baseAddress!.pointee)) {
+        let duration = TimeInterval(dispatchTimeInterval: startTime.distance(to: .now()))!
+        durationLabel.text = String(format: "%.3f seconds", duration)
+
+        let bufferPointer = UnsafeBufferPointer(start: output, count: labels.count)
+        let indexedResults = bufferPointer.enumerated()
+        let top3 = indexedResults.sorted { $0.1 > $1.1 }.prefix(4)
+
+        for result in top3 {
+          predictions.append(Prediction(confidence: result.1, label: String(labels[result.0])))
         }
-        resultView.text = text
+        resultsView.reloadData()
+      }
+    }
+  }
+
+  @IBAction func changeImage(_ sender: Any) {
+    updateImageView()
+    doPrediction()
+  }
+
+  @IBAction func runPredict(_ sender: Any) {
+    doPrediction()
+  }
+}
+
+extension ViewController {
+  func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
+    predictions.count
+  }
+
+  func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
+    let cell = tableView.dequeueReusableCell(withIdentifier: "result", for: indexPath)
+    cell.textLabel?.text = predictions[indexPath.row].label
+    cell.detailTextLabel?.text = String(format: "%.2f", predictions[indexPath.row].confidence)
+    return cell
+  }
+}
+
+extension TimeInterval {
+
+  init?(dispatchTimeInterval: DispatchTimeInterval) {
+    switch dispatchTimeInterval {
+    case .seconds(let value): self = Double(value)
+    case .milliseconds(let value): self = Double(value) / 1_000
+    case .microseconds(let value): self = Double(value) / 1_000_000
+    case .nanoseconds(let value):  self = Double(value) / 1_000_000_000
+    case .never: return nil
+    @unknown default: return nil
     }
+  }
 }
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-1.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-1.jpg
new file mode 100644
index 0000000..790afbc
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-1.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-2.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-2.jpg
new file mode 100644
index 0000000..3426175
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-2.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-3.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-3.jpg
new file mode 100644
index 0000000..7b2bd05
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/animal-3.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-1.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-1.jpg
new file mode 100644
index 0000000..c6dc1e6
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-1.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-2.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-2.jpg
new file mode 100644
index 0000000..b62d31d
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-2.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-3.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-3.jpg
new file mode 100644
index 0000000..4561529
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/car-3.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/food-1.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/food-1.jpg
new file mode 100644
index 0000000..2fb4e06
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/food-1.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-1.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-1.jpg
new file mode 100644
index 0000000..4fa80b6
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-1.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-2.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-2.jpg
new file mode 100644
index 0000000..6f0fc52
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-2.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-3.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-3.jpg
new file mode 100644
index 0000000..00c4ae0
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/girl-3.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/kitten_small.jpg b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/kitten_small.jpg
new file mode 100644
index 0000000..1592710
Binary files /dev/null and b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/kitten_small.jpg differ
diff --git a/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/image.png b/HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/wolf.png
similarity index 100%
rename from HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/image.png
rename to HelloWorld-Metal/HelloWorld-Metal/HelloWorld-Metal/images/wolf.png
diff --git a/HelloWorld/HelloWorld/HelloWorld.xcodeproj/project.pbxproj b/HelloWorld/HelloWorld/HelloWorld.xcodeproj/project.pbxproj
index 0e93380..f1901e7 100644
--- a/HelloWorld/HelloWorld/HelloWorld.xcodeproj/project.pbxproj
+++ b/HelloWorld/HelloWorld/HelloWorld.xcodeproj/project.pbxproj
@@ -7,7 +7,7 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
-		A0497B682352FFA000D70E9E /* image.png in Resources */ = {isa = PBXBuildFile; fileRef = A0497B672352FFA000D70E9E /* image.png */; };
+		0E9C55F7809AF58493258E5E /* libPods-HelloWorld.a in Frameworks */ = {isa = PBXBuildFile; fileRef = FDD5B9DA48EB1C10AB0378E1 /* libPods-HelloWorld.a */; };
 		A07753C62343DDEE00186969 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = A07753C52343DDEE00186969 /* AppDelegate.swift */; };
 		A07753CA2343DDEE00186969 /* ViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = A07753C92343DDEE00186969 /* ViewController.swift */; };
 		A07753CD2343DDEE00186969 /* Main.storyboard in Resources */ = {isa = PBXBuildFile; fileRef = A07753CB2343DDEE00186969 /* Main.storyboard */; };
@@ -17,10 +17,12 @@
 		A07753E32343F9AD00186969 /* model.pt in Resources */ = {isa = PBXBuildFile; fileRef = A07753E12343F9AD00186969 /* model.pt */; };
 		A07753E42343F9AD00186969 /* words.txt in Resources */ = {isa = PBXBuildFile; fileRef = A07753E22343F9AD00186969 /* words.txt */; };
 		A07753EA2343FDEC00186969 /* UIImage+Helper.swift in Sources */ = {isa = PBXBuildFile; fileRef = A07753E92343FDEC00186969 /* UIImage+Helper.swift */; };
+		BDCC15D32757ABAF005A979D /* images in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15D22757ABAE005A979D /* images */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXFileReference section */
-		A0497B672352FFA000D70E9E /* image.png */ = {isa = PBXFileReference; lastKnownFileType = image.png; path = image.png; sourceTree = "<group>"; };
+		1B8CD501236A9FDE6602CC95 /* Pods-HelloWorld.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-HelloWorld.release.xcconfig"; path = "Target Support Files/Pods-HelloWorld/Pods-HelloWorld.release.xcconfig"; sourceTree = "<group>"; };
+		7367FACB8B6E8FFF126590B1 /* Pods-HelloWorld.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-HelloWorld.debug.xcconfig"; path = "Target Support Files/Pods-HelloWorld/Pods-HelloWorld.debug.xcconfig"; sourceTree = "<group>"; };
 		A07753C22343DDEE00186969 /* HelloWorld.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = HelloWorld.app; sourceTree = BUILT_PRODUCTS_DIR; };
 		A07753C52343DDEE00186969 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = "<group>"; };
 		A07753C92343DDEE00186969 /* ViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ViewController.swift; sourceTree = "<group>"; };
@@ -34,6 +36,8 @@
 		A07753E12343F9AD00186969 /* model.pt */ = {isa = PBXFileReference; lastKnownFileType = file; path = model.pt; sourceTree = "<group>"; };
 		A07753E22343F9AD00186969 /* words.txt */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text; path = words.txt; sourceTree = "<group>"; };
 		A07753E92343FDEC00186969 /* UIImage+Helper.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = "UIImage+Helper.swift"; sourceTree = "<group>"; };
+		BDCC15D22757ABAE005A979D /* images */ = {isa = PBXFileReference; lastKnownFileType = folder; path = images; sourceTree = "<group>"; };
+		FDD5B9DA48EB1C10AB0378E1 /* libPods-HelloWorld.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-HelloWorld.a"; sourceTree = BUILT_PRODUCTS_DIR; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -41,6 +45,7 @@
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				0E9C55F7809AF58493258E5E /* libPods-HelloWorld.a in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -52,6 +57,8 @@
 			children = (
 				A07753C42343DDEE00186969 /* HelloWorld */,
 				A07753C32343DDEE00186969 /* Products */,
+				C13266EB510C04784F97934C /* Pods */,
+				BE1BE2724AE47886C898562B /* Frameworks */,
 			);
 			sourceTree = "<group>";
 		};
@@ -66,7 +73,7 @@
 		A07753C42343DDEE00186969 /* HelloWorld */ = {
 			isa = PBXGroup;
 			children = (
-				A0497B672352FFA000D70E9E /* image.png */,
+				BDCC15D22757ABAE005A979D /* images */,
 				A07753E02343F9AD00186969 /* model */,
 				A07753D92343EDE700186969 /* TorchBridge */,
 				A07753C52343DDEE00186969 /* AppDelegate.swift */,
@@ -99,6 +106,23 @@
 			path = model;
 			sourceTree = "<group>";
 		};
+		BE1BE2724AE47886C898562B /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				FDD5B9DA48EB1C10AB0378E1 /* libPods-HelloWorld.a */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
+		C13266EB510C04784F97934C /* Pods */ = {
+			isa = PBXGroup;
+			children = (
+				7367FACB8B6E8FFF126590B1 /* Pods-HelloWorld.debug.xcconfig */,
+				1B8CD501236A9FDE6602CC95 /* Pods-HelloWorld.release.xcconfig */,
+			);
+			path = Pods;
+			sourceTree = "<group>";
+		};
 /* End PBXGroup section */
 
 /* Begin PBXNativeTarget section */
@@ -106,6 +130,7 @@
 			isa = PBXNativeTarget;
 			buildConfigurationList = A07753D62343DDEF00186969 /* Build configuration list for PBXNativeTarget "HelloWorld" */;
 			buildPhases = (
+				5F2C0E985D616F4B61860016 /* [CP] Check Pods Manifest.lock */,
 				A07753BE2343DDEE00186969 /* Sources */,
 				A07753BF2343DDEE00186969 /* Frameworks */,
 				A07753C02343DDEE00186969 /* Resources */,
@@ -126,7 +151,7 @@
 			isa = PBXProject;
 			attributes = {
 				LastSwiftUpdateCheck = 1100;
-				LastUpgradeCheck = 1100;
+				LastUpgradeCheck = 1310;
 				ORGANIZATIONNAME = "Tao Xu";
 				TargetAttributes = {
 					A07753C12343DDEE00186969 = {
@@ -158,17 +183,42 @@
 			isa = PBXResourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				BDCC15D32757ABAF005A979D /* images in Resources */,
 				A07753D22343DDEF00186969 /* LaunchScreen.storyboard in Resources */,
 				A07753CF2343DDEF00186969 /* Assets.xcassets in Resources */,
 				A07753CD2343DDEE00186969 /* Main.storyboard in Resources */,
 				A07753E42343F9AD00186969 /* words.txt in Resources */,
 				A07753E32343F9AD00186969 /* model.pt in Resources */,
-				A0497B682352FFA000D70E9E /* image.png in Resources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
 /* End PBXResourcesBuildPhase section */
 
+/* Begin PBXShellScriptBuildPhase section */
+		5F2C0E985D616F4B61860016 /* [CP] Check Pods Manifest.lock */ = {
+			isa = PBXShellScriptBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			inputFileListPaths = (
+			);
+			inputPaths = (
+				"${PODS_PODFILE_DIR_PATH}/Podfile.lock",
+				"${PODS_ROOT}/Manifest.lock",
+			);
+			name = "[CP] Check Pods Manifest.lock";
+			outputFileListPaths = (
+			);
+			outputPaths = (
+				"$(DERIVED_FILE_DIR)/Pods-HelloWorld-checkManifestLockResult.txt",
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+			shellPath = /bin/sh;
+			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
+			showEnvVarsInLog = 0;
+		};
+/* End PBXShellScriptBuildPhase section */
+
 /* Begin PBXSourcesBuildPhase section */
 		A07753BE2343DDEE00186969 /* Sources */ = {
 			isa = PBXSourcesBuildPhase;
@@ -229,6 +279,7 @@
 				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
 				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
 				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
 				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
 				CLANG_WARN_STRICT_PROTOTYPES = YES;
 				CLANG_WARN_SUSPICIOUS_MOVE = YES;
@@ -289,6 +340,7 @@
 				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
 				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
 				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
 				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
 				CLANG_WARN_STRICT_PROTOTYPES = YES;
 				CLANG_WARN_SUSPICIOUS_MOVE = YES;
@@ -319,13 +371,15 @@
 		};
 		A07753D72343DDEF00186969 /* Debug */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = 7367FACB8B6E8FFF126590B1 /* Pods-HelloWorld.debug.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_STYLE = Automatic;
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = HelloWorld/Info.plist;
-				IPHONEOS_DEPLOYMENT_TARGET = 12.0;
+				IPHONEOS_DEPLOYMENT_TARGET = 14.0;
 				LD_RUNPATH_SEARCH_PATHS = (
 					"$(inherited)",
 					"@executable_path/Frameworks",
@@ -335,19 +389,21 @@
 				SWIFT_OBJC_BRIDGING_HEADER = "HelloWorld/TorchBridge/HelloWorld-Bridging-Header.h";
 				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
 				SWIFT_VERSION = 5.0;
-				TARGETED_DEVICE_FAMILY = 1;
+				TARGETED_DEVICE_FAMILY = "1,2";
 			};
 			name = Debug;
 		};
 		A07753D82343DDEF00186969 /* Release */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = 1B8CD501236A9FDE6602CC95 /* Pods-HelloWorld.release.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_STYLE = Automatic;
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = HelloWorld/Info.plist;
-				IPHONEOS_DEPLOYMENT_TARGET = 12.0;
+				IPHONEOS_DEPLOYMENT_TARGET = 14.0;
 				LD_RUNPATH_SEARCH_PATHS = (
 					"$(inherited)",
 					"@executable_path/Frameworks",
@@ -356,7 +412,7 @@
 				PRODUCT_NAME = "$(TARGET_NAME)";
 				SWIFT_OBJC_BRIDGING_HEADER = "HelloWorld/TorchBridge/HelloWorld-Bridging-Header.h";
 				SWIFT_VERSION = 5.0;
-				TARGETED_DEVICE_FAMILY = 1;
+				TARGETED_DEVICE_FAMILY = "1,2";
 			};
 			name = Release;
 		};
diff --git a/HelloWorld/HelloWorld/HelloWorld/Base.lproj/Main.storyboard b/HelloWorld/HelloWorld/HelloWorld/Base.lproj/Main.storyboard
index 21ead02..689298d 100644
--- a/HelloWorld/HelloWorld/HelloWorld/Base.lproj/Main.storyboard
+++ b/HelloWorld/HelloWorld/HelloWorld/Base.lproj/Main.storyboard
@@ -1,73 +1,133 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="14868" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="FxE-BR-UmL">
-    <device id="retina6_1" orientation="portrait" appearance="light"/>
-    <dependencies>
-        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="14824"/>
-        <capability name="Safe area layout guides" minToolsVersion="9.0"/>
-        <capability name="documents saved in the Xcode 8 format" minToolsVersion="8.0"/>
-    </dependencies>
-    <scenes>
-        <!--Image Classification-->
-        <scene sceneID="tne-QT-ifu">
-            <objects>
-                <viewController id="BYZ-38-t0r" customClass="ViewController" customModule="HelloWorld" customModuleProvider="target" sceneMemberID="viewController">
-                    <view key="view" contentMode="scaleToFill" id="8bC-Xf-vdC">
-                        <rect key="frame" x="0.0" y="0.0" width="414" height="896"/>
-                        <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
-                        <subviews>
-                            <imageView clipsSubviews="YES" userInteractionEnabled="NO" contentMode="scaleAspectFit" horizontalHuggingPriority="251" verticalHuggingPriority="251" translatesAutoresizingMaskIntoConstraints="NO" id="uCO-1m-OuG">
-                                <rect key="frame" x="20" y="108" width="374" height="374"/>
-                                <constraints>
-                                    <constraint firstAttribute="width" secondItem="uCO-1m-OuG" secondAttribute="height" multiplier="1:1" id="dDM-Le-7Fb"/>
-                                </constraints>
-                            </imageView>
-                            <textView clipsSubviews="YES" multipleTouchEnabled="YES" userInteractionEnabled="NO" contentMode="scaleToFill" scrollEnabled="NO" showsHorizontalScrollIndicator="NO" showsVerticalScrollIndicator="NO" editable="NO" selectable="NO" translatesAutoresizingMaskIntoConstraints="NO" id="Vo0-Cj-MiF">
-                                <rect key="frame" x="20" y="502" width="374" height="340"/>
-                                <color key="backgroundColor" systemColor="systemBackgroundColor" cocoaTouchSystemColor="whiteColor"/>
-                                <color key="textColor" systemColor="labelColor" cocoaTouchSystemColor="darkTextColor"/>
-                                <fontDescription key="fontDescription" type="system" pointSize="18"/>
-                                <textInputTraits key="textInputTraits" autocapitalizationType="sentences"/>
-                            </textView>
-                        </subviews>
-                        <color key="backgroundColor" systemColor="systemBackgroundColor" cocoaTouchSystemColor="whiteColor"/>
-                        <constraints>
-                            <constraint firstItem="uCO-1m-OuG" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="C6G-p9-aph"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="uCO-1m-OuG" secondAttribute="trailing" constant="20" id="JNn-cR-8v1"/>
-                            <constraint firstItem="uCO-1m-OuG" firstAttribute="top" secondItem="6Tk-OE-BBY" secondAttribute="top" constant="20" id="KrI-h7-nhY"/>
-                            <constraint firstItem="Vo0-Cj-MiF" firstAttribute="top" secondItem="uCO-1m-OuG" secondAttribute="bottom" constant="20" id="PaK-0f-UzX"/>
-                            <constraint firstItem="Vo0-Cj-MiF" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="S0H-11-FVR"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="bottom" secondItem="Vo0-Cj-MiF" secondAttribute="bottom" constant="20" id="djB-rZ-Vew"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="Vo0-Cj-MiF" secondAttribute="trailing" constant="20" id="qNP-eM-p0g"/>
-                        </constraints>
-                        <viewLayoutGuide key="safeArea" id="6Tk-OE-BBY"/>
-                    </view>
-                    <navigationItem key="navigationItem" title="Image Classification" id="din-cH-j77"/>
-                    <connections>
-                        <outlet property="imageView" destination="uCO-1m-OuG" id="gkN-Z5-OVB"/>
-                        <outlet property="resultView" destination="Vo0-Cj-MiF" id="lRB-Kj-nUX"/>
-                    </connections>
-                </viewController>
-                <placeholder placeholderIdentifier="IBFirstResponder" id="dkx-z0-nzr" sceneMemberID="firstResponder"/>
-            </objects>
-            <point key="canvasLocation" x="1042.0289855072465" y="137.94642857142856"/>
-        </scene>
-        <!--Navigation Controller-->
-        <scene sceneID="JQc-Lc-3UY">
-            <objects>
-                <navigationController automaticallyAdjustsScrollViewInsets="NO" id="FxE-BR-UmL" sceneMemberID="viewController">
-                    <toolbarItems/>
-                    <navigationBar key="navigationBar" contentMode="scaleToFill" insetsLayoutMarginsFromSafeArea="NO" id="7IF-Ml-aLx">
-                        <rect key="frame" x="0.0" y="44" width="414" height="44"/>
-                        <autoresizingMask key="autoresizingMask"/>
-                    </navigationBar>
-                    <nil name="viewControllers"/>
-                    <connections>
-                        <segue destination="BYZ-38-t0r" kind="relationship" relationship="rootViewController" id="1xO-jg-OSU"/>
-                    </connections>
-                </navigationController>
-                <placeholder placeholderIdentifier="IBFirstResponder" id="EbU-uZ-XIZ" userLabel="First Responder" customClass="UIResponder" sceneMemberID="firstResponder"/>
-            </objects>
-            <point key="canvasLocation" x="131.8840579710145" y="137.94642857142856"/>
-        </scene>
-    </scenes>
+<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="19455" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="MJX-KQ-XfU">
+<device id="retina6_1" orientation="portrait" appearance="light"/>
+<dependencies>
+<deployment identifier="iOS"/>
+<plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="19454"/>
+<capability name="Safe area layout guides" minToolsVersion="9.0"/>
+<capability name="System colors in document resources" minToolsVersion="11.0"/>
+<capability name="documents saved in the Xcode 8 format" minToolsVersion="8.0"/>
+</dependencies>
+<scenes>
+<!--Navigation Controller-->
+<scene sceneID="7wD-05-xk2">
+<objects>
+<navigationController id="MJX-KQ-XfU" sceneMemberID="viewController">
+<navigationBar key="navigationBar" contentMode="scaleToFill" id="Aq9-5f-Vyl">
+<rect key="frame" x="0.0" y="44" width="414" height="44"/>
+<autoresizingMask key="autoresizingMask"/>
+</navigationBar>
+<connections>
+<segue destination="BYZ-38-t0r" kind="relationship" relationship="rootViewController" id="ug1-NB-9L6"/>
+</connections>
+</navigationController>
+<placeholder placeholderIdentifier="IBFirstResponder" id="y4g-kY-87G" userLabel="First Responder" customClass="UIResponder" sceneMemberID="firstResponder"/>
+</objects>
+<point key="canvasLocation" x="-1133" y="-86"/>
+</scene>
+<!--Image Classification-->
+<scene sceneID="tne-QT-ifu">
+<objects>
+<viewController id="BYZ-38-t0r" customClass="ViewController" customModule="HelloWorld_Metal" customModuleProvider="target" sceneMemberID="viewController">
+<view key="view" contentMode="scaleToFill" id="8bC-Xf-vdC">
+<rect key="frame" x="0.0" y="0.0" width="414" height="896"/>
+<autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
+<subviews>
+<imageView clipsSubviews="YES" userInteractionEnabled="NO" contentMode="scaleAspectFit" horizontalHuggingPriority="251" verticalHuggingPriority="251" translatesAutoresizingMaskIntoConstraints="NO" id="tcN-Mf-gVj">
+<rect key="frame" x="20" y="108" width="374" height="374"/>
+<constraints>
+<constraint firstAttribute="width" secondItem="tcN-Mf-gVj" secondAttribute="height" multiplier="1:1" id="9wt-9w-9ry"/>
+</constraints>
+</imageView>
+<button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="aAV-kg-ted">
+<rect key="frame" x="172" y="831" width="70.5" height="31"/>
+<state key="normal" title="Button"/>
+<buttonConfiguration key="configuration" style="plain" title="Predict"/>
+<connections>
+<action selector="runPredict:" destination="BYZ-38-t0r" eventType="touchUpInside" id="pXp-cH-wQb"/>
+</connections>
+</button>
+<stepper opaque="NO" contentMode="scaleToFill" horizontalHuggingPriority="750" verticalHuggingPriority="750" contentHorizontalAlignment="center" contentVerticalAlignment="center" autorepeat="NO" wraps="YES" maximumValue="100" translatesAutoresizingMaskIntoConstraints="NO" id="qy4-9O-Fqi">
+<rect key="frame" x="160" y="490" width="94" height="32"/>
+<connections>
+<action selector="changeImage:" destination="BYZ-38-t0r" eventType="valueChanged" id="QnP-Mj-Zyu"/>
+</connections>
+</stepper>
+<tableView clipsSubviews="YES" contentMode="scaleToFill" alwaysBounceVertical="YES" dataMode="prototypes" style="plain" separatorStyle="none" allowsSelection="NO" rowHeight="45" estimatedRowHeight="-1" sectionHeaderHeight="-1" estimatedSectionHeaderHeight="-1" sectionFooterHeight="-1" estimatedSectionFooterHeight="-1" translatesAutoresizingMaskIntoConstraints="NO" id="Pdd-lq-6jU">
+<rect key="frame" x="20" y="559" width="374" height="264"/>
+<color key="backgroundColor" systemColor="systemBackgroundColor"/>
+<prototypes>
+<tableViewCell clipsSubviews="YES" contentMode="scaleToFill" insetsLayoutMarginsFromSafeArea="NO" selectionStyle="blue" hidesAccessoryWhenEditing="NO" indentationLevel="1" indentationWidth="0.0" reuseIdentifier="result" textLabel="PCe-zs-Bvh" detailTextLabel="489-Jh-vGK" rowHeight="46" style="IBUITableViewCellStyleSubtitle" id="Igw-8i-eEj">
+<rect key="frame" x="0.0" y="44.5" width="374" height="46"/>
+<autoresizingMask key="autoresizingMask"/>
+<tableViewCellContentView key="contentView" opaque="NO" clipsSubviews="YES" multipleTouchEnabled="YES" contentMode="center" insetsLayoutMarginsFromSafeArea="NO" tableViewCell="Igw-8i-eEj" id="0a0-zl-Zyy">
+<rect key="frame" x="0.0" y="0.0" width="374" height="46"/>
+<autoresizingMask key="autoresizingMask"/>
+<subviews>
+<label opaque="NO" multipleTouchEnabled="YES" contentMode="left" insetsLayoutMarginsFromSafeArea="NO" text="Title" textAlignment="natural" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" id="PCe-zs-Bvh">
+<rect key="frame" x="20" y="6.5" width="28.5" height="17"/>
+<autoresizingMask key="autoresizingMask"/>
+<fontDescription key="fontDescription" style="UICTFontTextStyleBody"/>
+<nil key="textColor"/>
+<nil key="highlightedColor"/>
+</label>
+<label opaque="NO" multipleTouchEnabled="YES" contentMode="left" insetsLayoutMarginsFromSafeArea="NO" text="Subtitle" textAlignment="natural" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" id="489-Jh-vGK">
+<rect key="frame" x="20" y="25" width="44" height="14.5"/>
+<autoresizingMask key="autoresizingMask"/>
+<fontDescription key="fontDescription" style="UICTFontTextStyleFootnote"/>
+<color key="textColor" systemColor="secondaryLabelColor"/>
+<nil key="highlightedColor"/>
+</label>
+</subviews>
+</tableViewCellContentView>
+</tableViewCell>
+</prototypes>
+</tableView>
+<label opaque="NO" userInteractionEnabled="NO" contentMode="left" horizontalHuggingPriority="251" verticalHuggingPriority="251" text="Label" textAlignment="natural" lineBreakMode="tailTruncation" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" translatesAutoresizingMaskIntoConstraints="NO" id="gbA-Ug-toR">
+<rect key="frame" x="186" y="530" width="42" height="21"/>
+<fontDescription key="fontDescription" type="system" pointSize="17"/>
+<nil key="textColor"/>
+<nil key="highlightedColor"/>
+</label>
+</subviews>
+<viewLayoutGuide key="safeArea" id="6Tk-OE-BBY"/>
+<color key="backgroundColor" systemColor="systemBackgroundColor"/>
+<constraints>
+<constraint firstItem="Pdd-lq-6jU" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="3rY-Q5-vyN"/>
+<constraint firstItem="aAV-kg-ted" firstAttribute="bottom" secondItem="6Tk-OE-BBY" secondAttribute="bottom" id="4Yg-He-1ZE"/>
+<constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="tcN-Mf-gVj" secondAttribute="trailing" constant="20" id="6qF-bk-R5a"/>
+<constraint firstItem="gbA-Ug-toR" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="8ae-F0-4Zh"/>
+<constraint firstItem="qy4-9O-Fqi" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="9wT-68-2w9"/>
+<constraint firstItem="tcN-Mf-gVj" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="20" id="Gig-Ia-Utu"/>
+<constraint firstItem="aAV-kg-ted" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="Pb2-2G-usQ"/>
+<constraint firstItem="gbA-Ug-toR" firstAttribute="top" secondItem="qy4-9O-Fqi" secondAttribute="bottom" constant="8" symbolic="YES" id="eya-SQ-EHq"/>
+<constraint firstItem="aAV-kg-ted" firstAttribute="top" secondItem="Pdd-lq-6jU" secondAttribute="bottom" constant="8" symbolic="YES" id="lBU-d2-eRA"/>
+<constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="Pdd-lq-6jU" secondAttribute="trailing" constant="20" id="lpI-D5-kgi"/>
+<constraint firstItem="qy4-9O-Fqi" firstAttribute="top" secondItem="tcN-Mf-gVj" secondAttribute="bottom" constant="8" symbolic="YES" id="nZz-51-QLt"/>
+<constraint firstItem="tcN-Mf-gVj" firstAttribute="top" secondItem="6Tk-OE-BBY" secondAttribute="top" constant="20" id="uBS-lr-amS"/>
+<constraint firstItem="Pdd-lq-6jU" firstAttribute="top" secondItem="gbA-Ug-toR" secondAttribute="bottom" constant="8" symbolic="YES" id="xJv-jJ-eEP"/>
+</constraints>
+</view>
+<navigationItem key="navigationItem" title="Image Classification" id="n4K-Na-VKn"/>
+<connections>
+<outlet property="durationLabel" destination="gbA-Ug-toR" id="zvH-Xo-NdQ"/>
+<outlet property="imageStepper" destination="qy4-9O-Fqi" id="CRJ-Sz-qW8"/>
+<outlet property="imageView" destination="tcN-Mf-gVj" id="9y4-1U-aPc"/>
+<outlet property="predictButton" destination="aAV-kg-ted" id="beT-pE-GIZ"/>
+<outlet property="resultsView" destination="Pdd-lq-6jU" id="wDn-Fy-cI8"/>
+</connections>
+</viewController>
+<placeholder placeholderIdentifier="IBFirstResponder" id="dkx-z0-nzr" sceneMemberID="firstResponder"/>
+</objects>
+<point key="canvasLocation" x="-278.26086956521743" y="-86.383928571428569"/>
+</scene>
+</scenes>
+<resources>
+<systemColor name="secondaryLabelColor">
+<color red="0.23529411764705882" green="0.23529411764705882" blue="0.2627450980392157" alpha="0.59999999999999998" colorSpace="custom" customColorSpace="sRGB"/>
+</systemColor>
+<systemColor name="systemBackgroundColor">
+<color white="1" alpha="1" colorSpace="custom" customColorSpace="genericGamma22GrayColorSpace"/>
+</systemColor>
+</resources>
 </document>
diff --git a/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.h b/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.h
index 883120f..78b339b 100644
--- a/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.h
+++ b/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.h
@@ -4,11 +4,11 @@ NS_ASSUME_NONNULL_BEGIN
 
 @interface TorchModule : NSObject
 
-- (nullable instancetype)initWithFileAtPath:(NSString*)filePath
-    NS_SWIFT_NAME(init(fileAtPath:))NS_DESIGNATED_INITIALIZER;
+- (nullable instancetype)initWithFileAtPath:(NSString*)filePath width:(long)width height:(long)height
+NS_SWIFT_NAME(init(fileAtPath:width:height:))NS_DESIGNATED_INITIALIZER;
 + (instancetype)new NS_UNAVAILABLE;
 - (instancetype)init NS_UNAVAILABLE;
-- (nullable NSArray<NSNumber*>*)predictImage:(void*)imageBuffer NS_SWIFT_NAME(predict(image:));
+- (nullable float*)predictImage:(void*)imageBuffer NS_SWIFT_NAME(predict(image:));
 
 @end
 
diff --git a/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.mm b/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.mm
index 3f6155c..02a9318 100644
--- a/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.mm
+++ b/HelloWorld/HelloWorld/HelloWorld/TorchBridge/TorchModule.mm
@@ -1,14 +1,20 @@
 #import "TorchModule.h"
-#import <Libtorch-Lite/Libtorch-Lite.h>
+#import "LibTorch-Lite.h"
 
 @implementation TorchModule {
- @protected
+@protected
   torch::jit::mobile::Module _impl;
+  at::Tensor _output;
+  int64_t _inputSize[4];
 }
 
-- (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
+- (nullable instancetype)initWithFileAtPath:(NSString*)filePath width:(long)width height:(long)height {
   self = [super init];
   if (self) {
+    _inputSize[0] = 1;
+    _inputSize[1] = 3;
+    _inputSize[2] = width;
+    _inputSize[3] = height;
     try {
       _impl = torch::jit::_load_for_mobile(filePath.UTF8String);
     } catch (const std::exception& exception) {
@@ -19,23 +25,16 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
   return self;
 }
 
-- (NSArray<NSNumber*>*)predictImage:(void*)imageBuffer {
+- (float*)predictImage:(void*)imageBuffer {
   try {
-    at::Tensor tensor = torch::from_blob(imageBuffer, {1, 3, 224, 224}, at::kFloat);
-    c10::InferenceMode guard;
-    auto outputTensor = _impl.forward({tensor}).toTensor();
-    float* floatBuffer = outputTensor.data_ptr<float>();
-    if (!floatBuffer) {
-      return nil;
-    }
-    NSMutableArray* results = [[NSMutableArray alloc] init];
-    for (int i = 0; i < 1000; i++) {
-      [results addObject:@(floatBuffer[i])];
-    }
-    return [results copy];
+    c10::InferenceMode guard(true);
+    at::Tensor tensor = torch::from_blob(imageBuffer, _inputSize, at::kFloat);
+    _output = _impl.forward({tensor}).toTensor().cpu();
+    return _output.data_ptr<float>();
   } catch (const std::exception& exception) {
     NSLog(@"%s", exception.what());
   }
+
   return nil;
 }
 
diff --git a/HelloWorld/HelloWorld/HelloWorld/ViewController.swift b/HelloWorld/HelloWorld/HelloWorld/ViewController.swift
index 96894eb..2092d88 100644
--- a/HelloWorld/HelloWorld/HelloWorld/ViewController.swift
+++ b/HelloWorld/HelloWorld/HelloWorld/ViewController.swift
@@ -1,43 +1,124 @@
 import UIKit
 
-class ViewController: UIViewController {
-    @IBOutlet var imageView: UIImageView!
-    @IBOutlet var resultView: UITextView!
-    private lazy var module: TorchModule = {
-        if let filePath = Bundle.main.path(forResource: "model", ofType: "pt"),
-            let module = TorchModule(fileAtPath: filePath) {
-            return module
-        } else {
-            fatalError("Can't find the model file!")
-        }
-    }()
-
-    private lazy var labels: [String] = {
-        if let filePath = Bundle.main.path(forResource: "words", ofType: "txt"),
-            let labels = try? String(contentsOfFile: filePath) {
-            return labels.components(separatedBy: .newlines)
-        } else {
-            fatalError("Can't find the text file!")
-        }
-    }()
-
-    override func viewDidLoad() {
-        super.viewDidLoad()
-        let image = UIImage(named: "image.png")!
-        imageView.image = image
-        let resizedImage = image.resized(to: CGSize(width: 224, height: 224))
-        guard var pixelBuffer = resizedImage.normalized() else {
-            return
-        }
-        guard let outputs = module.predict(image: UnsafeMutableRawPointer(&pixelBuffer)) else {
-            return
-        }
-        let zippedResults = zip(labels.indices, outputs)
-        let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(3)
-        var text = ""
-        for result in sortedResults {
-            text += "\u{2022} \(labels[result.0]) \n\n"
+struct Prediction {
+  let confidence: Float
+  let label: String
+}
+
+class ViewController: UIViewController, UITableViewDataSource, UITableViewDelegate {
+  @IBOutlet weak var imageView: UIImageView!
+  @IBOutlet weak var predictButton: UIButton!
+  @IBOutlet weak var imageStepper: UIStepper!
+  @IBOutlet weak var durationLabel: UILabel!
+  @IBOutlet weak var resultsView: UITableView!
+
+  var predictions = [Prediction]()
+
+  let width: Int = 224
+  let height: Int = 224
+
+  private lazy var module: TorchModule = {
+    guard let filePath = Bundle.main.path(forResource: "model", ofType: "pt"),
+          let module = TorchModule(fileAtPath: filePath, width: width, height: height)
+    else {
+      fatalError("Can't find the model file!")
+    }
+    return module
+  }()
+
+  private lazy var labels: [Substring] = {
+    guard let filePath = Bundle.main.path(forResource: "words", ofType: "txt"),
+          let labels = try? String(contentsOfFile: filePath)
+    else {
+      fatalError("Can't find the text file!")
+    }
+    return labels.split(whereSeparator: \.isNewline)
+  }()
+
+  private lazy var images: [UIImage?] = {
+    let imagePaths = Bundle.main.paths(forResourcesOfType: "jpg", inDirectory: "images") +
+    Bundle.main.paths(forResourcesOfType: "png", inDirectory: "images")
+    print(imagePaths)
+    return imagePaths.map { UIImage(named: $0) }
+  }()
+
+  private var activeImage: UIImage? { images[Int(imageStepper.value)] }
+
+  override func viewDidLoad() {
+    super.viewDidLoad()
+    resultsView.dataSource = self
+    imageStepper.value = 0
+    imageStepper.minimumValue = 0
+    imageStepper.maximumValue = Double(images.count - 1)
+    updateImageView()
+    doPrediction()
+  }
+
+  func updateImageView() {
+    imageView.image = images[Int(imageStepper.value)]
+  }
+
+  private var imageFloats: [Float32]? {
+    guard let resizedImage = activeImage?.resized(to: CGSize(width: width, height: height)) else { return nil }
+    return resizedImage.normalized()
+  }
+
+  func doPrediction() {
+    guard var input = self.imageFloats else { return }
+    predictions.removeAll()
+    resultsView.reloadData()
+
+    input.withUnsafeMutableBufferPointer { ptr in
+      let startTime: DispatchTime = .now()
+      if let output = module.predict(image: &(ptr.baseAddress!.pointee)) {
+        let duration = TimeInterval(dispatchTimeInterval: startTime.distance(to: .now()))!
+        durationLabel.text = String(format: "%.3f seconds", duration)
+
+        let bufferPointer = UnsafeBufferPointer(start: output, count: labels.count)
+        let indexedResults = bufferPointer.enumerated()
+        let top3 = indexedResults.sorted { $0.1 > $1.1 }.prefix(4)
+
+        for result in top3 {
+          predictions.append(Prediction(confidence: result.1, label: String(labels[result.0])))
         }
-        resultView.text = text
+        resultsView.reloadData()
+      }
+    }
+  }
+
+  @IBAction func changeImage(_ sender: Any) {
+    updateImageView()
+    doPrediction()
+  }
+
+  @IBAction func runPredict(_ sender: Any) {
+    doPrediction()
+  }
+}
+
+extension ViewController {
+  func tableView(_ tableView: UITableView, numberOfRowsInSection section: Int) -> Int {
+    predictions.count
+  }
+
+  func tableView(_ tableView: UITableView, cellForRowAt indexPath: IndexPath) -> UITableViewCell {
+    let cell = tableView.dequeueReusableCell(withIdentifier: "result", for: indexPath)
+    cell.textLabel?.text = predictions[indexPath.row].label
+    cell.detailTextLabel?.text = String(format: "%.2f", predictions[indexPath.row].confidence)
+    return cell
+  }
+}
+
+extension TimeInterval {
+
+  init?(dispatchTimeInterval: DispatchTimeInterval) {
+    switch dispatchTimeInterval {
+    case .seconds(let value): self = Double(value)
+    case .milliseconds(let value): self = Double(value) / 1_000
+    case .microseconds(let value): self = Double(value) / 1_000_000
+    case .nanoseconds(let value):  self = Double(value) / 1_000_000_000
+    case .never: return nil
+    @unknown default: return nil
     }
+  }
 }
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/animal-1.jpg b/HelloWorld/HelloWorld/HelloWorld/images/animal-1.jpg
new file mode 100644
index 0000000..790afbc
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/animal-1.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/animal-2.jpg b/HelloWorld/HelloWorld/HelloWorld/images/animal-2.jpg
new file mode 100644
index 0000000..3426175
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/animal-2.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/animal-3.jpg b/HelloWorld/HelloWorld/HelloWorld/images/animal-3.jpg
new file mode 100644
index 0000000..7b2bd05
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/animal-3.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/car-1.jpg b/HelloWorld/HelloWorld/HelloWorld/images/car-1.jpg
new file mode 100644
index 0000000..c6dc1e6
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/car-1.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/car-2.jpg b/HelloWorld/HelloWorld/HelloWorld/images/car-2.jpg
new file mode 100644
index 0000000..b62d31d
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/car-2.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/car-3.jpg b/HelloWorld/HelloWorld/HelloWorld/images/car-3.jpg
new file mode 100644
index 0000000..4561529
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/car-3.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/food-1.jpg b/HelloWorld/HelloWorld/HelloWorld/images/food-1.jpg
new file mode 100644
index 0000000..2fb4e06
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/food-1.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/girl-1.jpg b/HelloWorld/HelloWorld/HelloWorld/images/girl-1.jpg
new file mode 100644
index 0000000..4fa80b6
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/girl-1.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/girl-2.jpg b/HelloWorld/HelloWorld/HelloWorld/images/girl-2.jpg
new file mode 100644
index 0000000..6f0fc52
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/girl-2.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/girl-3.jpg b/HelloWorld/HelloWorld/HelloWorld/images/girl-3.jpg
new file mode 100644
index 0000000..00c4ae0
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/girl-3.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/images/kitten_small.jpg b/HelloWorld/HelloWorld/HelloWorld/images/kitten_small.jpg
new file mode 100644
index 0000000..1592710
Binary files /dev/null and b/HelloWorld/HelloWorld/HelloWorld/images/kitten_small.jpg differ
diff --git a/HelloWorld/HelloWorld/HelloWorld/image.png b/HelloWorld/HelloWorld/HelloWorld/images/wolf.png
similarity index 100%
rename from HelloWorld/HelloWorld/HelloWorld/image.png
rename to HelloWorld/HelloWorld/HelloWorld/images/wolf.png
diff --git a/PyTorchDemo/Podfile b/PyTorchDemo/Podfile
index 5a0aeef..50504d0 100644
--- a/PyTorchDemo/Podfile
+++ b/PyTorchDemo/Podfile
@@ -1,4 +1,5 @@
 platform :ios, '14.3'
 target 'PyTorchDemo' do
-  pod 'LibTorch', '~>1.9.0'
+  # pod 'LibTorch', '~>1.9.0'
+  pod 'LibTorch-Lite-Nightly'
 end
diff --git a/PyTorchDemo/PyTorchDemo.xcodeproj/project.pbxproj b/PyTorchDemo/PyTorchDemo.xcodeproj/project.pbxproj
index 08e6725..26041cd 100644
--- a/PyTorchDemo/PyTorchDemo.xcodeproj/project.pbxproj
+++ b/PyTorchDemo/PyTorchDemo.xcodeproj/project.pbxproj
@@ -7,6 +7,7 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		6D650F9479F0EB0B7DB8EA95 /* libPods-PyTorchDemo.a in Frameworks */ = {isa = PBXBuildFile; fileRef = 6022D6FD02DB575F4D75312F /* libPods-PyTorchDemo.a */; };
 		A01F52AC233C6C32004039FC /* topics.txt in Resources */ = {isa = PBXBuildFile; fileRef = A01F52AB233C6C32004039FC /* topics.txt */; };
 		A01F52AE233C985C004039FC /* Predictor.swift in Sources */ = {isa = PBXBuildFile; fileRef = A01F52AD233C985C004039FC /* Predictor.swift */; };
 		A0296C4A2333178500C77366 /* NLPViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = A0296C492333178500C77366 /* NLPViewController.swift */; };
@@ -41,9 +42,11 @@
 		A09E1796233990A0002F8F15 /* VisionModelCard.xib in Resources */ = {isa = PBXBuildFile; fileRef = A09E1795233990A0002F8F15 /* VisionModelCard.xib */; };
 		A0B12D4E2338944900380B1D /* TorchModule.mm in Sources */ = {isa = PBXBuildFile; fileRef = A0B12D4D2338944900380B1D /* TorchModule.mm */; };
 		B1C35BA725C99B5200CAB118 /* trace_model.py in Resources */ = {isa = PBXBuildFile; fileRef = B1C35BA625C99B5200CAB118 /* trace_model.py */; };
+		BDCC15D52757DCD6005A979D /* mobilenet_quantized2.pt in Resources */ = {isa = PBXBuildFile; fileRef = BDCC15D42757DCD6005A979D /* mobilenet_quantized2.pt */; };
 /* End PBXBuildFile section */
 
 /* Begin PBXFileReference section */
+		6022D6FD02DB575F4D75312F /* libPods-PyTorchDemo.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = "libPods-PyTorchDemo.a"; sourceTree = BUILT_PRODUCTS_DIR; };
 		A01F52AB233C6C32004039FC /* topics.txt */ = {isa = PBXFileReference; lastKnownFileType = text; path = topics.txt; sourceTree = "<group>"; };
 		A01F52AD233C985C004039FC /* Predictor.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Predictor.swift; sourceTree = "<group>"; };
 		A0296C492333178500C77366 /* NLPViewController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = NLPViewController.swift; sourceTree = "<group>"; };
@@ -82,6 +85,9 @@
 		A0B12D4C2338944900380B1D /* TorchModule.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = TorchModule.h; sourceTree = "<group>"; };
 		A0B12D4D2338944900380B1D /* TorchModule.mm */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.cpp.objcpp; path = TorchModule.mm; sourceTree = "<group>"; };
 		B1C35BA625C99B5200CAB118 /* trace_model.py */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = text.script.python; path = trace_model.py; sourceTree = "<group>"; };
+		BDCC15D42757DCD6005A979D /* mobilenet_quantized2.pt */ = {isa = PBXFileReference; lastKnownFileType = file; path = mobilenet_quantized2.pt; sourceTree = "<group>"; };
+		CE2E63B3E0FE1515FBF41C12 /* Pods-PyTorchDemo.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-PyTorchDemo.debug.xcconfig"; path = "Target Support Files/Pods-PyTorchDemo/Pods-PyTorchDemo.debug.xcconfig"; sourceTree = "<group>"; };
+		FA2C77CD1130895F363C81F0 /* Pods-PyTorchDemo.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-PyTorchDemo.release.xcconfig"; path = "Target Support Files/Pods-PyTorchDemo/Pods-PyTorchDemo.release.xcconfig"; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -89,6 +95,7 @@
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				6D650F9479F0EB0B7DB8EA95 /* libPods-PyTorchDemo.a in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -143,6 +150,8 @@
 			children = (
 				A091B19523303F2D0008A5B8 /* PyTorchDemo */,
 				A091B19423303F2D0008A5B8 /* Products */,
+				FB40E40FC092817D5F1A6B47 /* Pods */,
+				F39C5143BC568706EC1AA82A /* Frameworks */,
 			);
 			sourceTree = "<group>";
 		};
@@ -190,6 +199,7 @@
 			isa = PBXGroup;
 			children = (
 				B1C35BA625C99B5200CAB118 /* trace_model.py */,
+				BDCC15D42757DCD6005A979D /* mobilenet_quantized2.pt */,
 				A065C92A2335711B006BEAAD /* words.txt */,
 				A03349BB233D5E7C003425FD /* mobilenet_quantized.pt */,
 			);
@@ -219,6 +229,23 @@
 			path = TorchBridge;
 			sourceTree = "<group>";
 		};
+		F39C5143BC568706EC1AA82A /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				6022D6FD02DB575F4D75312F /* libPods-PyTorchDemo.a */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
+		FB40E40FC092817D5F1A6B47 /* Pods */ = {
+			isa = PBXGroup;
+			children = (
+				CE2E63B3E0FE1515FBF41C12 /* Pods-PyTorchDemo.debug.xcconfig */,
+				FA2C77CD1130895F363C81F0 /* Pods-PyTorchDemo.release.xcconfig */,
+			);
+			path = Pods;
+			sourceTree = "<group>";
+		};
 /* End PBXGroup section */
 
 /* Begin PBXNativeTarget section */
@@ -226,6 +253,7 @@
 			isa = PBXNativeTarget;
 			buildConfigurationList = A091B1A523303F2D0008A5B8 /* Build configuration list for PBXNativeTarget "PyTorchDemo" */;
 			buildPhases = (
+				83A344EDF737CBA77AA96B94 /* [CP] Check Pods Manifest.lock */,
 				A091B18F23303F2D0008A5B8 /* Sources */,
 				A091B19023303F2D0008A5B8 /* Frameworks */,
 				A091B19123303F2D0008A5B8 /* Resources */,
@@ -246,7 +274,7 @@
 			isa = PBXProject;
 			attributes = {
 				LastSwiftUpdateCheck = 1030;
-				LastUpgradeCheck = 1030;
+				LastUpgradeCheck = 1310;
 				TargetAttributes = {
 					A091B19223303F2D0008A5B8 = {
 						CreatedOnToolsVersion = 10.3;
@@ -284,6 +312,7 @@
 				A091B19E23303F2D0008A5B8 /* Assets.xcassets in Resources */,
 				A03349CA233DAE9F003425FD /* NLPResultView.xib in Resources */,
 				A03349C4233D71F2003425FD /* ImageClassificationItemView.xib in Resources */,
+				BDCC15D52757DCD6005A979D /* mobilenet_quantized2.pt in Resources */,
 				A036EA33233C19F500054A36 /* reddit.pt in Resources */,
 				B1C35BA725C99B5200CAB118 /* trace_model.py in Resources */,
 				A03349CE233DB128003425FD /* NLPItemView.xib in Resources */,
@@ -296,6 +325,31 @@
 		};
 /* End PBXResourcesBuildPhase section */
 
+/* Begin PBXShellScriptBuildPhase section */
+		83A344EDF737CBA77AA96B94 /* [CP] Check Pods Manifest.lock */ = {
+			isa = PBXShellScriptBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			inputFileListPaths = (
+			);
+			inputPaths = (
+				"${PODS_PODFILE_DIR_PATH}/Podfile.lock",
+				"${PODS_ROOT}/Manifest.lock",
+			);
+			name = "[CP] Check Pods Manifest.lock";
+			outputFileListPaths = (
+			);
+			outputPaths = (
+				"$(DERIVED_FILE_DIR)/Pods-PyTorchDemo-checkManifestLockResult.txt",
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+			shellPath = /bin/sh;
+			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
+			showEnvVarsInLog = 0;
+		};
+/* End PBXShellScriptBuildPhase section */
+
 /* Begin PBXSourcesBuildPhase section */
 		A091B18F23303F2D0008A5B8 /* Sources */ = {
 			isa = PBXSourcesBuildPhase;
@@ -372,6 +426,7 @@
 				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
 				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
 				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
 				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
 				CLANG_WARN_STRICT_PROTOTYPES = YES;
 				CLANG_WARN_SUSPICIOUS_MOVE = YES;
@@ -433,6 +488,7 @@
 				CLANG_WARN_OBJC_IMPLICIT_RETAIN_SELF = YES;
 				CLANG_WARN_OBJC_LITERAL_CONVERSION = YES;
 				CLANG_WARN_OBJC_ROOT_CLASS = YES_ERROR;
+				CLANG_WARN_QUOTED_INCLUDE_IN_FRAMEWORK_HEADER = YES;
 				CLANG_WARN_RANGE_LOOP_ANALYSIS = YES;
 				CLANG_WARN_STRICT_PROTOTYPES = YES;
 				CLANG_WARN_SUSPICIOUS_MOVE = YES;
@@ -464,12 +520,13 @@
 		};
 		A091B1A623303F2D0008A5B8 /* Debug */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = CE2E63B3E0FE1515FBF41C12 /* Pods-PyTorchDemo.debug.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_IDENTITY = "iPhone Developer";
 				CODE_SIGN_STYLE = Automatic;
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = PyTorchDemo/Info.plist;
 				LD_RUNPATH_SEARCH_PATHS = (
@@ -482,18 +539,19 @@
 				SWIFT_OBJC_BRIDGING_HEADER = "PyTorchDemo/TorchBridge/PyTorchDemo-Bridging-Header.h";
 				SWIFT_OPTIMIZATION_LEVEL = "-Onone";
 				SWIFT_VERSION = 5.0;
-				TARGETED_DEVICE_FAMILY = 1;
+				TARGETED_DEVICE_FAMILY = "1,2";
 			};
 			name = Debug;
 		};
 		A091B1A723303F2D0008A5B8 /* Release */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = FA2C77CD1130895F363C81F0 /* Pods-PyTorchDemo.release.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_IDENTITY = "iPhone Developer";
 				CODE_SIGN_STYLE = Automatic;
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = PyTorchDemo/Info.plist;
 				LD_RUNPATH_SEARCH_PATHS = (
@@ -505,7 +563,7 @@
 				PROVISIONING_PROFILE_SPECIFIER = "";
 				SWIFT_OBJC_BRIDGING_HEADER = "PyTorchDemo/TorchBridge/PyTorchDemo-Bridging-Header.h";
 				SWIFT_VERSION = 5.0;
-				TARGETED_DEVICE_FAMILY = 1;
+				TARGETED_DEVICE_FAMILY = "1,2";
 			};
 			name = Release;
 		};
diff --git a/PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift b/PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift
index f232cf4..639b326 100644
--- a/PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift
+++ b/PyTorchDemo/PyTorchDemo/ImageClassification/ImagePredictor.swift
@@ -1,38 +1,40 @@
 import UIKit
 
 class ImagePredictor: Predictor {
-    private var isRunning: Bool = false
-    private lazy var module: VisionTorchModule = {
-        if let filePath = Bundle.main.path(forResource: "mobilenet_quantized", ofType: "pt"),
-            let module = VisionTorchModule(fileAtPath: filePath) {
-            return module
-        } else {
-            fatalError("Failed to load model!")
-        }
-    }()
+  private var isRunning: Bool = false
 
-    private var labels: [String] = {
-        if let filePath = Bundle.main.path(forResource: "words", ofType: "txt"),
-            let labels = try? String(contentsOfFile: filePath) {
-            return labels.components(separatedBy: .newlines)
-        } else {
-            fatalError("Label file was not found.")
-        }
-    }()
+  private lazy var module: VisionTorchModule = {
+    guard let filePath = Bundle.main.path(forResource: "mobilenet_quantized2", ofType: "pt"),
+       let module = VisionTorchModule(fileAtPath: filePath)
+    else {
+      fatalError("Failed to load model!")
+    }
+    return module
+  }()
+
+  private var labels: [String] = {
+    guard let filePath = Bundle.main.path(forResource: "words", ofType: "txt"),
+          let labels = try? String(contentsOfFile: filePath)
+    else {
+      fatalError("Label file was not found.")
+    }
+    return labels.split(whereSeparator: \.isNewline).map { .init($0) }
+  }()
+
+  func predict(_ buffer: [Float32], resultCount: Int) throws -> ([InferenceResult], Double)? {
+    guard !isRunning else { return nil }
+    isRunning = true
+    defer { isRunning = false }
 
-    func predict(_ buffer: [Float32], resultCount: Int) throws -> ([InferenceResult], Double)? {
-        if isRunning {
-            return nil
-        }
-        isRunning = true
-        let startTime = CACurrentMediaTime()
-        var tensorBuffer = buffer;
-        guard let outputs = module.predict(image: UnsafeMutableRawPointer(&tensorBuffer)) else {
-            throw PredictorError.invalidInputTensor
-        }
-        isRunning = false
-        let inferenceTime = (CACurrentMediaTime() - startTime) * 1000
-        let results = topK(scores: outputs, labels: labels, count: resultCount)
-        return (results, inferenceTime)
+    var tmp = buffer
+    return tmp.withUnsafeMutableBufferPointer { ptr in
+      let startTime = CACurrentMediaTime()
+      if let output = module.predict(image: &(ptr.baseAddress!.pointee)) {
+        let duration = (CACurrentMediaTime() - startTime) * 1000
+        let bufferPointer = UnsafeBufferPointer(start: output, count: labels.count)
+        return (topK(scores: bufferPointer, labels: labels, count: 3), duration)
+      }
+      return nil
     }
+  }
 }
diff --git a/PyTorchDemo/PyTorchDemo/ImageClassification/model/mobilenet_quantized2.pt b/PyTorchDemo/PyTorchDemo/ImageClassification/model/mobilenet_quantized2.pt
new file mode 100644
index 0000000..ad4c5cd
Binary files /dev/null and b/PyTorchDemo/PyTorchDemo/ImageClassification/model/mobilenet_quantized2.pt differ
diff --git a/PyTorchDemo/PyTorchDemo/ImageClassification/model/trace_model.py b/PyTorchDemo/PyTorchDemo/ImageClassification/model/trace_model.py
index 6aa0084..eeb02fc 100644
--- a/PyTorchDemo/PyTorchDemo/ImageClassification/model/trace_model.py
+++ b/PyTorchDemo/PyTorchDemo/ImageClassification/model/trace_model.py
@@ -3,8 +3,7 @@
 from torch.utils.mobile_optimizer import optimize_for_mobile
 
 model = torchvision.models.quantization.mobilenet_v2(pretrained=True, quantize=True)
-model.eval()
 example = torch.rand(1, 3, 224, 224)
 traced_script_module = torch.jit.trace(model, example)
 torchscript_model_optimized = optimize_for_mobile(traced_script_module)
-torchscript_model_optimized.save("mobilenet_quantized.pt")
+torchscript_model_optimized._save_for_lite_interpreter("mobilenet_quantized2.pt")
diff --git a/PyTorchDemo/PyTorchDemo/NLP/NLPPredictor.swift b/PyTorchDemo/PyTorchDemo/NLP/NLPPredictor.swift
index 68a482d..6bd7573 100644
--- a/PyTorchDemo/PyTorchDemo/NLP/NLPPredictor.swift
+++ b/PyTorchDemo/PyTorchDemo/NLP/NLPPredictor.swift
@@ -1,34 +1,35 @@
 import UIKit
 
 class NLPPredictor: Predictor {
-    private var module: NLPTorchModule = {
-        if let filePath = Bundle.main.path(forResource: "reddit", ofType: "pt"),
-            let module = NLPTorchModule(fileAtPath: filePath) {
-            return module
-        } else {
-            fatalError("Failed to load model!")
-        }
-    }()
-
-    private var topics: [String] = []
-    init() {
-        topics = loadTopics()
+  private var module: NLPTorchModule = {
+    if let filePath = Bundle.main.path(forResource: "reddit", ofType: "pt"),
+       let module = NLPTorchModule(fileAtPath: filePath) {
+      return module
+    } else {
+      fatalError("Failed to load model!")
     }
+  }()
+
+  private var topics: [String] = []
 
-    func predict(_ text: String, resultCount: Int) throws -> [InferenceResult]? {
-        if text.isEmpty {
-            throw PredictorError.invalidInputTensor
-        }
-        guard let outputs = module.predict(text: text) else {
-            throw PredictorError.invalidInputTensor
-        }
-        return topK(scores: outputs, labels: topics, count: resultCount)
+  init() {
+    topics = loadTopics()
+  }
+
+  func predict(_ text: String, resultCount: Int) throws -> [InferenceResult]? {
+    guard !text.isEmpty,
+          let output = module.predict(text: text) else {
+      throw PredictorError.invalidInputTensor
     }
 
-    private func loadTopics() -> [String] {
-        guard let topics = module.topics() else {
-            fatalError("Failed to load topics from model")
-        }
-        return topics
+    let floatArray = UnsafeBufferPointer(start: output, count: resultCount)
+    return topK(scores: floatArray, labels: topics, count: resultCount)
+  }
+
+  private func loadTopics() -> [String] {
+    guard let topics = module.topics() else {
+      fatalError("Failed to load topics from model")
     }
+    return topics
+  }
 }
diff --git a/PyTorchDemo/PyTorchDemo/Predictor.swift b/PyTorchDemo/PyTorchDemo/Predictor.swift
index 3ad7351..98dde4d 100644
--- a/PyTorchDemo/PyTorchDemo/Predictor.swift
+++ b/PyTorchDemo/PyTorchDemo/Predictor.swift
@@ -14,9 +14,9 @@ enum PredictorError: Swift.Error {
 protocol Predictor {}
 
 extension Predictor {
-    func topK(scores: [NSNumber], labels: [String], count: Int) -> [InferenceResult] {
-        let zippedResults = zip(labels.indices, scores)
-        let sortedResults = zippedResults.sorted { $0.1.floatValue > $1.1.floatValue }.prefix(count)
-        return sortedResults.map { InferenceResult(score: $0.1.floatValue, label: labels[$0.0]) }
-    }
+  func topK(scores: UnsafeBufferPointer<Float>, labels: [String], count: Int) -> [InferenceResult] {
+    let zippedResults = zip(scores, labels.indices)
+    let sortedResults = zippedResults.sorted { $0.0 > $1.0 }.prefix(count)
+    return sortedResults.map { InferenceResult(score: $0.0, label: labels[$0.1]) }
+  }
 }
diff --git a/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.h b/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.h
index f4435d3..e14efed 100644
--- a/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.h
+++ b/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.h
@@ -12,12 +12,12 @@ NS_ASSUME_NONNULL_BEGIN
 @end
 
 @interface VisionTorchModule : TorchModule
-- (nullable NSArray<NSNumber*>*)predictImage:(void*)imageBuffer NS_SWIFT_NAME(predict(image:));
+- (nullable float*)predictImage:(void*)imageBuffer NS_SWIFT_NAME(predict(image:));
 @end
 
 @interface NLPTorchModule : TorchModule
 - (nullable NSArray<NSString*>*)topics;
-- (nullable NSArray<NSNumber*>*)predictText:(NSString*)text NS_SWIFT_NAME(predict(text:));
+- (nullable float*)predictText:(NSString*)text NS_SWIFT_NAME(predict(text:));
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.mm b/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.mm
index b188794..035bf24 100644
--- a/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.mm
+++ b/PyTorchDemo/PyTorchDemo/TorchBridge/TorchModule.mm
@@ -1,21 +1,17 @@
 #import "TorchModule.h"
-#import <LibTorch/LibTorch.h>
+#import <LibTorch-Lite.h>
 
 @implementation TorchModule {
- @protected
-  torch::jit::script::Module _impl;
+@protected
+  torch::jit::mobile::Module _model;
+  at::Tensor _output;
 }
 
 - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
   self = [super init];
   if (self) {
     try {
-      auto qengines = at::globalContext().supportedQEngines();
-      if (std::find(qengines.begin(), qengines.end(), at::QEngine::QNNPACK) != qengines.end()) {
-        at::globalContext().setQEngine(at::QEngine::QNNPACK);
-      }
-      _impl = torch::jit::load(filePath.UTF8String);
-      _impl.eval();
+      _model = torch::jit::_load_for_mobile(filePath.UTF8String);
     } catch (const std::exception& exception) {
       NSLog(@"%s", exception.what());
       return nil;
@@ -28,21 +24,12 @@ - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
 
 @implementation VisionTorchModule
 
-- (NSArray<NSNumber*>*)predictImage:(void*)imageBuffer {
+- (float*)predictImage:(void*)imageBuffer {
   try {
+    c10::InferenceMode guard(true);
     at::Tensor tensor = torch::from_blob(imageBuffer, {1, 3, 224, 224}, at::kFloat);
-    torch::autograd::AutoGradMode guard(false);
-    at::AutoNonVariableTypeMode non_var_type_mode(true);
-    auto outputTensor = _impl.forward({tensor}).toTensor();
-    float* floatBuffer = outputTensor.data_ptr<float>();
-    if (!floatBuffer) {
-      return nil;
-    }
-    NSMutableArray* results = [[NSMutableArray alloc] init];
-    for (int i = 0; i < 1000; i++) {
-      [results addObject:@(floatBuffer[i])];
-    }
-    return [results copy];
+    _output = _model.forward({tensor}).toTensor();
+    return _output.data_ptr<float>();
   } catch (const std::exception& exception) {
     NSLog(@"%s", exception.what());
   }
@@ -53,22 +40,13 @@ @implementation VisionTorchModule
 
 @implementation NLPTorchModule
 
-- (NSArray<NSNumber*>*)predictText:(NSString*)text {
+- (float*)predictText:(NSString*)text {
   try {
+    c10::InferenceMode guard(true);
     const char* buffer = text.UTF8String;
-    torch::autograd::AutoGradMode guard(false);
-    at::AutoNonVariableTypeMode non_var_type_mode(true);
     at::Tensor tensor = torch::from_blob((void*)buffer, {1, (int64_t)(strlen(buffer))}, at::kByte);
-    auto outputTensor = _impl.forward({tensor}).toTensor();
-    float* floatBuffer = outputTensor.data_ptr<float>();
-    if (!floatBuffer) {
-      return nil;
-    }
-    NSMutableArray* results = [[NSMutableArray alloc] init];
-    for (int i = 0; i < 16; i++) {
-      [results addObject:@(floatBuffer[i])];
-    }
-    return [results copy];
+    _output = _model.forward({tensor}).toTensor();
+    return _output.data_ptr<float>();
   } catch (const std::exception& exception) {
     NSLog(@"%s", exception.what());
   }
@@ -77,7 +55,7 @@ @implementation NLPTorchModule
 
 - (NSArray<NSString*>*)topics {
   try {
-    auto genericList = _impl.run_method("get_classes").toList();
+    auto genericList = _model.run_method("get_classes").toList();
     NSMutableArray<NSString*>* topics = [NSMutableArray<NSString*> new];
     for (int i = 0; i < genericList.size(); i++) {
       std::string topic = genericList.get(i).toString()->string();
diff --git a/TorchVideo/TorchVideo.xcodeproj/project.pbxproj b/TorchVideo/TorchVideo.xcodeproj/project.pbxproj
index 7d62381..a905c56 100644
--- a/TorchVideo/TorchVideo.xcodeproj/project.pbxproj
+++ b/TorchVideo/TorchVideo.xcodeproj/project.pbxproj
@@ -7,6 +7,7 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		150FEDD29415E812ECDA20C1 /* Pods_TorchVideo.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 630BF504711D114B0B8F709D /* Pods_TorchVideo.framework */; };
 		26024B6426042B2300B6F625 /* AppDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26024B6326042B2300B6F625 /* AppDelegate.swift */; };
 		26024B6626042B2300B6F625 /* SceneDelegate.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26024B6526042B2300B6F625 /* SceneDelegate.swift */; };
 		26024B6826042B2300B6F625 /* ViewController.swift in Sources */ = {isa = PBXBuildFile; fileRef = 26024B6726042B2300B6F625 /* ViewController.swift */; };
@@ -29,6 +30,7 @@
 /* End PBXBuildFile section */
 
 /* Begin PBXFileReference section */
+		06D85729F77BC3AA53A0C8C4 /* Pods-TorchVideo.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-TorchVideo.debug.xcconfig"; path = "Target Support Files/Pods-TorchVideo/Pods-TorchVideo.debug.xcconfig"; sourceTree = "<group>"; };
 		26024B6026042B2300B6F625 /* TorchVideo.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = TorchVideo.app; sourceTree = BUILT_PRODUCTS_DIR; };
 		26024B6326042B2300B6F625 /* AppDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppDelegate.swift; sourceTree = "<group>"; };
 		26024B6526042B2300B6F625 /* SceneDelegate.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SceneDelegate.swift; sourceTree = "<group>"; };
@@ -52,6 +54,8 @@
 		26024BA6260D0E8A00B6F625 /* CameraPreviewView.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = CameraPreviewView.swift; path = TorchVideo/CameraPreviewView.swift; sourceTree = SOURCE_ROOT; };
 		26024BA7260D0E8A00B6F625 /* CVPixelBuffer+Helper.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; name = "CVPixelBuffer+Helper.swift"; path = "TorchVideo/CVPixelBuffer+Helper.swift"; sourceTree = SOURCE_ROOT; };
 		26A8C10F26E17A3B00F4A58D /* video_classification.ptl */ = {isa = PBXFileReference; lastKnownFileType = file; path = video_classification.ptl; sourceTree = "<group>"; };
+		630BF504711D114B0B8F709D /* Pods_TorchVideo.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Pods_TorchVideo.framework; sourceTree = BUILT_PRODUCTS_DIR; };
+		729E5C95BA2DF8C99A950320 /* Pods-TorchVideo.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-TorchVideo.release.xcconfig"; path = "Target Support Files/Pods-TorchVideo/Pods-TorchVideo.release.xcconfig"; sourceTree = "<group>"; };
 /* End PBXFileReference section */
 
 /* Begin PBXFrameworksBuildPhase section */
@@ -59,6 +63,7 @@
 			isa = PBXFrameworksBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				150FEDD29415E812ECDA20C1 /* Pods_TorchVideo.framework in Frameworks */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;
 		};
@@ -71,6 +76,7 @@
 				26024B6226042B2300B6F625 /* TorchVideo */,
 				26024B6126042B2300B6F625 /* Products */,
 				EA0C558862607C5A9BA15E19 /* Pods */,
+				50FD079E11052576D34D17AB /* Frameworks */,
 			);
 			sourceTree = "<group>";
 		};
@@ -135,9 +141,19 @@
 			name = Utils;
 			sourceTree = "<group>";
 		};
+		50FD079E11052576D34D17AB /* Frameworks */ = {
+			isa = PBXGroup;
+			children = (
+				630BF504711D114B0B8F709D /* Pods_TorchVideo.framework */,
+			);
+			name = Frameworks;
+			sourceTree = "<group>";
+		};
 		EA0C558862607C5A9BA15E19 /* Pods */ = {
 			isa = PBXGroup;
 			children = (
+				06D85729F77BC3AA53A0C8C4 /* Pods-TorchVideo.debug.xcconfig */,
+				729E5C95BA2DF8C99A950320 /* Pods-TorchVideo.release.xcconfig */,
 			);
 			path = Pods;
 			sourceTree = "<group>";
@@ -149,6 +165,7 @@
 			isa = PBXNativeTarget;
 			buildConfigurationList = 26024B7426042B2500B6F625 /* Build configuration list for PBXNativeTarget "TorchVideo" */;
 			buildPhases = (
+				B0C22985D2F2F00E69387BC5 /* [CP] Check Pods Manifest.lock */,
 				26024B5C26042B2300B6F625 /* Sources */,
 				26024B5D26042B2300B6F625 /* Frameworks */,
 				26024B5E26042B2300B6F625 /* Resources */,
@@ -213,6 +230,31 @@
 		};
 /* End PBXResourcesBuildPhase section */
 
+/* Begin PBXShellScriptBuildPhase section */
+		B0C22985D2F2F00E69387BC5 /* [CP] Check Pods Manifest.lock */ = {
+			isa = PBXShellScriptBuildPhase;
+			buildActionMask = 2147483647;
+			files = (
+			);
+			inputFileListPaths = (
+			);
+			inputPaths = (
+				"${PODS_PODFILE_DIR_PATH}/Podfile.lock",
+				"${PODS_ROOT}/Manifest.lock",
+			);
+			name = "[CP] Check Pods Manifest.lock";
+			outputFileListPaths = (
+			);
+			outputPaths = (
+				"$(DERIVED_FILE_DIR)/Pods-TorchVideo-checkManifestLockResult.txt",
+			);
+			runOnlyForDeploymentPostprocessing = 0;
+			shellPath = /bin/sh;
+			shellScript = "diff \"${PODS_PODFILE_DIR_PATH}/Podfile.lock\" \"${PODS_ROOT}/Manifest.lock\" > /dev/null\nif [ $? != 0 ] ; then\n    # print error to STDERR\n    echo \"error: The sandbox is not in sync with the Podfile.lock. Run 'pod install' or update your CocoaPods installation.\" >&2\n    exit 1\nfi\n# This output is used by Xcode 'outputs' to avoid re-running this script phase.\necho \"SUCCESS\" > \"${SCRIPT_OUTPUT_FILE_0}\"\n";
+			showEnvVarsInLog = 0;
+		};
+/* End PBXShellScriptBuildPhase section */
+
 /* Begin PBXSourcesBuildPhase section */
 		26024B5C26042B2300B6F625 /* Sources */ = {
 			isa = PBXSourcesBuildPhase;
@@ -372,12 +414,13 @@
 		};
 		26024B7526042B2500B6F625 /* Debug */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = 06D85729F77BC3AA53A0C8C4 /* Pods-TorchVideo.debug.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_STYLE = Automatic;
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = TorchVideo/Info.plist;
 				IPHONEOS_DEPLOYMENT_TARGET = 14.0;
@@ -396,12 +439,13 @@
 		};
 		26024B7626042B2500B6F625 /* Release */ = {
 			isa = XCBuildConfiguration;
+			baseConfigurationReference = 729E5C95BA2DF8C99A950320 /* Pods-TorchVideo.release.xcconfig */;
 			buildSettings = {
 				ASSETCATALOG_COMPILER_APPICON_NAME = AppIcon;
 				ASSETCATALOG_COMPILER_GLOBAL_ACCENT_COLOR_NAME = AccentColor;
 				CLANG_ENABLE_MODULES = YES;
 				CODE_SIGN_STYLE = Automatic;
-				DEVELOPMENT_TEAM = "";
+				DEVELOPMENT_TEAM = UP6SS5ES7E;
 				ENABLE_BITCODE = NO;
 				INFOPLIST_FILE = TorchVideo/Info.plist;
 				IPHONEOS_DEPLOYMENT_TARGET = 14.0;
diff --git a/TorchVideo/TorchVideo/AppDelegate.swift b/TorchVideo/TorchVideo/AppDelegate.swift
index e742747..4f92cd0 100644
--- a/TorchVideo/TorchVideo/AppDelegate.swift
+++ b/TorchVideo/TorchVideo/AppDelegate.swift
@@ -11,26 +11,24 @@ import UIKit
 @main
 class AppDelegate: UIResponder, UIApplicationDelegate {
 
-
-
-    func application(_ application: UIApplication, didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey: Any]?) -> Bool {
-        // Override point for customization after application launch.
-        return true
-    }
-
-    // MARK: UISceneSession Lifecycle
-
-    func application(_ application: UIApplication, configurationForConnecting connectingSceneSession: UISceneSession, options: UIScene.ConnectionOptions) -> UISceneConfiguration {
-        // Called when a new scene session is being created.
-        // Use this method to select a configuration to create the new scene with.
-        return UISceneConfiguration(name: "Default Configuration", sessionRole: connectingSceneSession.role)
-    }
-
-    func application(_ application: UIApplication, didDiscardSceneSessions sceneSessions: Set<UISceneSession>) {
-        // Called when the user discards a scene session.
-        // If any sessions were discarded while the application was not running, this will be called shortly after application:didFinishLaunchingWithOptions.
-        // Use this method to release any resources that were specific to the discarded scenes, as they will not return.
-    }
+  func application(_ application: UIApplication, didFinishLaunchingWithOptions launchOptions: [UIApplication.LaunchOptionsKey: Any]?) -> Bool {
+    // Override point for customization after application launch.
+    return true
+  }
+
+  // MARK: UISceneSession Lifecycle
+
+  func application(_ application: UIApplication, configurationForConnecting connectingSceneSession: UISceneSession, options: UIScene.ConnectionOptions) -> UISceneConfiguration {
+    // Called when a new scene session is being created.
+    // Use this method to select a configuration to create the new scene with.
+    return UISceneConfiguration(name: "Default Configuration", sessionRole: connectingSceneSession.role)
+  }
+
+  func application(_ application: UIApplication, didDiscardSceneSessions sceneSessions: Set<UISceneSession>) {
+    // Called when the user discards a scene session.
+    // If any sessions were discarded while the application was not running, this will be called shortly after application:didFinishLaunchingWithOptions.
+    // Use this method to release any resources that were specific to the discarded scenes, as they will not return.
+  }
 
 
 }
diff --git a/TorchVideo/TorchVideo/Base.lproj/Main.storyboard b/TorchVideo/TorchVideo/Base.lproj/Main.storyboard
index 1aa98b5..605e6ef 100644
--- a/TorchVideo/TorchVideo/Base.lproj/Main.storyboard
+++ b/TorchVideo/TorchVideo/Base.lproj/Main.storyboard
@@ -1,8 +1,9 @@
 <?xml version="1.0" encoding="UTF-8"?>
-<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="17701" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="1Pk-ic-63J">
+<document type="com.apple.InterfaceBuilder3.CocoaTouch.Storyboard.XIB" version="3.0" toolsVersion="19455" targetRuntime="iOS.CocoaTouch" propertyAccessControl="none" useAutolayout="YES" useTraitCollections="YES" useSafeAreas="YES" colorMatched="YES" initialViewController="1Pk-ic-63J">
     <device id="retina6_1" orientation="portrait" appearance="light"/>
     <dependencies>
-        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="17703"/>
+        <deployment identifier="iOS"/>
+        <plugIn identifier="com.apple.InterfaceBuilder.IBCocoaTouchPlugin" version="19454"/>
         <capability name="Safe area layout guides" minToolsVersion="9.0"/>
         <capability name="System colors in document resources" minToolsVersion="11.0"/>
         <capability name="documents saved in the Xcode 8 format" minToolsVersion="8.0"/>
@@ -16,49 +17,49 @@
                         <rect key="frame" x="0.0" y="0.0" width="414" height="896"/>
                         <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
                         <subviews>
-                            <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="BAh-gy-asI">
-                                <rect key="frame" x="64" y="537" width="43" height="41"/>
-                                <fontDescription key="fontDescription" type="system" pointSize="24"/>
-                                <state key="normal" title="Test"/>
-                                <connections>
-                                    <action selector="testTapped:" destination="BYZ-38-t0r" eventType="touchUpInside" id="5df-Cb-lUm"/>
-                                </connections>
-                            </button>
-                            <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="DQi-c5-Jqa">
-                                <rect key="frame" x="187" y="537" width="65" height="41"/>
-                                <fontDescription key="fontDescription" type="system" pointSize="24"/>
-                                <state key="normal" title="Select"/>
-                                <connections>
-                                    <action selector="selectTapped:" destination="BYZ-38-t0r" eventType="touchUpInside" id="O0a-Y0-zXr"/>
-                                </connections>
-                            </button>
-                            <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="MHq-hZ-5AA">
-                                <rect key="frame" x="292" y="537" width="43" height="41"/>
-                                <fontDescription key="fontDescription" type="system" pointSize="24"/>
-                                <state key="normal" title="Live"/>
-                                <connections>
-                                    <segue destination="mcG-bu-5pK" kind="show" id="jvI-ox-35y"/>
-                                </connections>
-                            </button>
-                            <label opaque="NO" userInteractionEnabled="NO" contentMode="left" horizontalHuggingPriority="251" verticalHuggingPriority="251" fixedFrame="YES" text="To play pre-recorded videos, tap Test
To select video from Camera Roll, tap Select
To see live video classification, tap Live" textAlignment="natural" lineBreakMode="tailTruncation" numberOfLines="3" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" translatesAutoresizingMaskIntoConstraints="NO" id="5gU-dN-k8N">
-                                <rect key="frame" x="64" y="594" width="271" height="105"/>
-                                <autoresizingMask key="autoresizingMask" flexibleMaxX="YES" flexibleMaxY="YES"/>
+                            <label opaque="NO" userInteractionEnabled="NO" contentMode="left" horizontalHuggingPriority="251" verticalHuggingPriority="251" text="To play pre-recorded videos, tap Test
To select video from Camera Roll, tap Select
To see live video classification, tap Live" textAlignment="natural" lineBreakMode="tailTruncation" numberOfLines="3" baselineAdjustment="alignBaselines" adjustsFontSizeToFit="NO" translatesAutoresizingMaskIntoConstraints="NO" id="5gU-dN-k8N">
+                                <rect key="frame" x="73" y="815" width="268" height="47"/>
                                 <fontDescription key="fontDescription" type="system" pointSize="13"/>
                                 <nil key="textColor"/>
                                 <nil key="highlightedColor"/>
                             </label>
+                            <stackView opaque="NO" contentMode="scaleToFill" spacing="60" translatesAutoresizingMaskIntoConstraints="NO" id="2AH-fH-Zes">
+                                <rect key="frame" x="71.5" y="766" width="271" height="41"/>
+                                <subviews>
+                                    <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="BAh-gy-asI">
+                                        <rect key="frame" x="0.0" y="0.0" width="43" height="41"/>
+                                        <fontDescription key="fontDescription" type="system" pointSize="24"/>
+                                        <state key="normal" title="Test"/>
+                                        <connections>
+                                            <action selector="testTapped:" destination="BYZ-38-t0r" eventType="touchUpInside" id="5df-Cb-lUm"/>
+                                        </connections>
+                                    </button>
+                                    <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="DQi-c5-Jqa">
+                                        <rect key="frame" x="103" y="0.0" width="65" height="41"/>
+                                        <fontDescription key="fontDescription" type="system" pointSize="24"/>
+                                        <state key="normal" title="Select"/>
+                                        <connections>
+                                            <action selector="selectTapped:" destination="BYZ-38-t0r" eventType="touchUpInside" id="O0a-Y0-zXr"/>
+                                        </connections>
+                                    </button>
+                                    <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="MHq-hZ-5AA">
+                                        <rect key="frame" x="228" y="0.0" width="43" height="41"/>
+                                        <fontDescription key="fontDescription" type="system" pointSize="24"/>
+                                        <state key="normal" title="Live"/>
+                                        <connections>
+                                            <segue destination="mcG-bu-5pK" kind="show" id="jvI-ox-35y"/>
+                                        </connections>
+                                    </button>
+                                </subviews>
+                            </stackView>
                         </subviews>
                         <viewLayoutGuide key="safeArea" id="6Tk-OE-BBY"/>
                         <color key="backgroundColor" systemColor="systemBackgroundColor"/>
                         <constraints>
-                            <constraint firstItem="DQi-c5-Jqa" firstAttribute="leading" secondItem="8bC-Xf-vdC" secondAttribute="leading" constant="187" id="45n-wf-2Yf"/>
-                            <constraint firstItem="MHq-hZ-5AA" firstAttribute="leading" relation="greaterThanOrEqual" secondItem="DQi-c5-Jqa" secondAttribute="trailing" constant="8" symbolic="YES" id="5P5-74-Mu4"/>
-                            <constraint firstItem="DQi-c5-Jqa" firstAttribute="baseline" secondItem="MHq-hZ-5AA" secondAttribute="baseline" id="9Lt-F5-F2s"/>
-                            <constraint firstItem="DQi-c5-Jqa" firstAttribute="baseline" secondItem="BAh-gy-asI" secondAttribute="baseline" id="H2o-ap-ykd"/>
-                            <constraint firstItem="DQi-c5-Jqa" firstAttribute="leading" relation="greaterThanOrEqual" secondItem="BAh-gy-asI" secondAttribute="trailing" constant="8" symbolic="YES" id="P3d-DL-kgn"/>
-                            <constraint firstItem="BAh-gy-asI" firstAttribute="leading" secondItem="6Tk-OE-BBY" secondAttribute="leading" constant="64" id="T1U-rQ-p5c"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="trailing" secondItem="MHq-hZ-5AA" secondAttribute="trailing" constant="79" id="YFI-L1-Xhp"/>
-                            <constraint firstItem="6Tk-OE-BBY" firstAttribute="bottom" secondItem="DQi-c5-Jqa" secondAttribute="bottom" constant="284" id="kJQ-9y-2BD"/>
+                            <constraint firstItem="2AH-fH-Zes" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="7Sh-b9-a20"/>
+                            <constraint firstItem="5gU-dN-k8N" firstAttribute="centerX" secondItem="8bC-Xf-vdC" secondAttribute="centerX" id="IZC-79-RGG"/>
+                            <constraint firstItem="5gU-dN-k8N" firstAttribute="bottom" secondItem="6Tk-OE-BBY" secondAttribute="bottom" id="WtL-4V-6gd"/>
+                            <constraint firstItem="5gU-dN-k8N" firstAttribute="top" secondItem="2AH-fH-Zes" secondAttribute="bottom" constant="8" symbolic="YES" id="uw5-hK-mDq"/>
                         </constraints>
                     </view>
                     <navigationItem key="navigationItem" id="k0a-rL-UWW"/>
@@ -78,7 +79,7 @@
                         <rect key="frame" x="0.0" y="0.0" width="414" height="896"/>
                         <autoresizingMask key="autoresizingMask" widthSizable="YES" heightSizable="YES"/>
                         <subviews>
-                            <view contentMode="scaleToFill" translatesAutoresizingMaskIntoConstraints="NO" id="83V-Un-WeH" customClass="CameraPreviewView" customModule="ObjectDetection" customModuleProvider="target">
+                            <view contentMode="scaleToFill" translatesAutoresizingMaskIntoConstraints="NO" id="83V-Un-WeH" customClass="CameraPreviewView" customModule="TorchVideo" customModuleProvider="target">
                                 <rect key="frame" x="0.0" y="118" width="414" height="630"/>
                                 <subviews>
                                     <button opaque="NO" contentMode="scaleToFill" contentHorizontalAlignment="center" contentVerticalAlignment="center" buttonType="system" lineBreakMode="middleTruncation" translatesAutoresizingMaskIntoConstraints="NO" id="7zP-Jr-oTE">
diff --git a/TorchVideo/TorchVideo/InferenceModule.h b/TorchVideo/TorchVideo/InferenceModule.h
index 075def9..3fa9ebf 100644
--- a/TorchVideo/TorchVideo/InferenceModule.h
+++ b/TorchVideo/TorchVideo/InferenceModule.h
@@ -14,7 +14,7 @@ NS_ASSUME_NONNULL_BEGIN
     NS_SWIFT_NAME(init(fileAtPath:))NS_DESIGNATED_INITIALIZER;
 + (instancetype)new NS_UNAVAILABLE;
 - (instancetype)init NS_UNAVAILABLE;
-- (nullable NSArray<NSNumber*>*)classifyFrames:(void*)framesBuffer NS_SWIFT_NAME(classify(frames:));
+- (nullable float*)classifyFrames:(void*)framesBuffer NS_SWIFT_NAME(classify(frames:));
 
 @end
 
diff --git a/TorchVideo/TorchVideo/InferenceModule.mm b/TorchVideo/TorchVideo/InferenceModule.mm
index 626b4c6..dc8e61f 100644
--- a/TorchVideo/TorchVideo/InferenceModule.mm
+++ b/TorchVideo/TorchVideo/InferenceModule.mm
@@ -14,78 +14,33 @@
 
 
 @implementation InferenceModule {
-    @protected torch::jit::mobile::Module _impl;
+@protected torch::jit::mobile::Module _model;
+@protected at::Tensor _output;
 }
 
-NSInteger arrayWithIndexSort(NSArray* first, NSArray* second, void* context)
-{
-    id firstValue = [first objectAtIndex:0];
-    id secondValue = [second objectAtIndex:0];
-    return [secondValue compare:firstValue];
-}
-
-
 - (nullable instancetype)initWithFileAtPath:(NSString*)filePath {
-    self = [super init];
-    if (self) {
-        try {
-            _impl = torch::jit::_load_for_mobile(filePath.UTF8String);
-        } catch (const std::exception& exception) {
-            NSLog(@"%s", exception.what());
-            return nil;
-        }
-    }
-    return self;
-}
-
-- (int)argMax:(NSArray*)array {
-    int maxIdx = 0;
-    float maxVal = -FLT_MAX;
-    for (int j = 0; j < OUTPUT_SIZE; j++) {
-      if ([array[j] floatValue]> maxVal) {
-          maxVal = [array[j] floatValue];
-          maxIdx = j;
-      }
-    }
-    return maxIdx;
-}
-
-
-- (NSArray<NSNumber*>*)classifyFrames:(void*)framesBuffer {
+  self = [super init];
+  if (self) {
     try {
-        at::Tensor tensor = torch::from_blob(framesBuffer, { 1, 3, 4, INPUT_WIDTH, INPUT_HEIGHT }, at::kFloat);
-        
-        c10::InferenceMode guard;
-        
-        auto outputTensor = _impl.forward({ tensor }).toTensor();
-
-        float* floatBuffer = outputTensor.data_ptr<float>();
-        if (!floatBuffer) {
-            return nil;
-        }
-        
-        NSMutableArray* scores = [[NSMutableArray alloc] init];
-        for (int i = 0; i < OUTPUT_SIZE; i++) {
-          [scores addObject:@(floatBuffer[i])];
-        }
-        
-        NSMutableArray* scoresIdx = [[NSMutableArray alloc] init];
-        for (int i = 0; i < OUTPUT_SIZE; i++) {
-            [scoresIdx addObject:[NSArray arrayWithObjects:scores[i], @(i), nil]];
-        }
-        
-        NSArray* sortedScoresIdx = [scoresIdx sortedArrayUsingFunction:arrayWithIndexSort context:NULL];
-        
-        NSMutableArray* results = [[NSMutableArray alloc] init];
-        for (int i = 0; i < TOP_COUNT; i++)
-            [results addObject: sortedScoresIdx[i][1]];
-        
-        return [results copy];
-        
+      _model = torch::jit::_load_for_mobile(filePath.UTF8String);
     } catch (const std::exception& exception) {
-        NSLog(@"%s", exception.what());
+      NSLog(@"%s", exception.what());
+      return nil;
     }
-    return nil;
+  }
+  return self;
+}
+
+- (float*)classifyFrames:(void*)framesBuffer {
+  try {
+    c10::InferenceMode guard;
+    at::Tensor tensor = torch::from_blob(framesBuffer, { 1, 3, 4, INPUT_WIDTH, INPUT_HEIGHT }, at::kFloat);
+    _output = _model.forward({ tensor }).toTensor();
+    return _output.data_ptr<float>();
+  } catch (const std::exception& exception) {
+    NSLog(@"%s", exception.what());
+  }
+  return nil;
 }
 
 @end
diff --git a/TorchVideo/TorchVideo/LiveVideoClassificationViewController.swift b/TorchVideo/TorchVideo/LiveVideoClassificationViewController.swift
index c036a85..65d31cc 100644
--- a/TorchVideo/TorchVideo/LiveVideoClassificationViewController.swift
+++ b/TorchVideo/TorchVideo/LiveVideoClassificationViewController.swift
@@ -8,64 +8,69 @@ import AVFoundation
 import UIKit
 
 class LiveVideoClassificationViewController: ViewController {
-    @IBOutlet var cameraView: CameraPreviewView!
-    @IBOutlet var lblResult: UILabel!
-    @IBOutlet var indicator: UIActivityIndicatorView!
+  @IBOutlet var cameraView: CameraPreviewView!
+  @IBOutlet var lblResult: UILabel!
+  @IBOutlet var indicator: UIActivityIndicatorView!
 
-    private let delayMs: Double = 1000
-    private var prevTimestampMs: Double = 0.0
-    private var cameraController = CameraController()
-    private var imageViewLive =  UIImageView()
-    private var inferencer = VideoClassifier()
-    
-    override func viewDidLoad() {
-        super.viewDidLoad()
-        self.navigationController?.setNavigationBarHidden(false, animated: true)
-        
-        cameraController.configPreviewLayer(cameraView)
-        imageViewLive.frame = CGRect(x: 0, y: 0, width: cameraView.frame.size.width, height: cameraView.frame.size.height)
-        cameraView.addSubview(imageViewLive)
-        
-        cameraController.videoCaptureCompletionBlock = { [weak self] buffer, error in
-            guard let strongSelf = self else { return }
-            if error != nil {
-                return
-            }
-            guard var pixelBuffer = buffer else { return }
-            
-            // simulate 4 frames of image as requested by model input
-            pixelBuffer += pixelBuffer
-            pixelBuffer += pixelBuffer
-            
-            let currentTimestamp = CACurrentMediaTime()
-            if (currentTimestamp - strongSelf.prevTimestampMs) * 1000 <= strongSelf.delayMs { return }
-            strongSelf.prevTimestampMs = currentTimestamp
-            let startTime = CACurrentMediaTime()
-            guard let top5Indexes = self?.inferencer.module.classify(frames: &pixelBuffer) else {
-                return
-            }
-            let inferenceTime = CACurrentMediaTime() - startTime
-                
-            DispatchQueue.main.async {
-                let results = top5Indexes.map { self!.inferencer.classes[$0.intValue] }
-                self!.lblResult.text = results.joined(separator: ", ") + " - \(Int(1000*inferenceTime))ms"
+  private let delayMs: Double = 1000
+  private var prevTimestampMs: Double = 0.0
+  private var cameraController = CameraController()
+  private var imageViewLive =  UIImageView()
+  private var inferencer = VideoClassifier()
 
-            }
-        }
-    }
+  override func viewDidLoad() {
+    super.viewDidLoad()
+    self.navigationController?.setNavigationBarHidden(false, animated: true)
 
-    override func viewWillAppear(_ animated: Bool) {
-        super.viewWillAppear(animated)
-        navigationController?.setNavigationBarHidden(true, animated: false)
-        cameraController.startSession()
-    }
+    cameraController.configPreviewLayer(cameraView)
+    imageViewLive.frame = CGRect(x: 0, y: 0, width: cameraView.frame.size.width, height: cameraView.frame.size.height)
+    cameraView.addSubview(imageViewLive)
 
-    override func viewWillDisappear(_ animated: Bool) {
-        super.viewWillDisappear(animated)
-        cameraController.stopSession()
-    }
+    cameraController.videoCaptureCompletionBlock = { [weak self] buffer, error in
+      let currentTimestamp = CACurrentMediaTime()
+      guard let strongSelf = self,
+            var pixelBuffer = buffer,
+            error == nil,
+            (currentTimestamp - strongSelf.prevTimestampMs) * 1000 > strongSelf.delayMs
+      else {
+        return
+      }
+
+      strongSelf.prevTimestampMs = currentTimestamp
+
+      // simulate 4 frames of image as requested by model input
+      pixelBuffer += pixelBuffer
+      pixelBuffer += pixelBuffer
 
-    @IBAction func onBackClicked(_: Any) {
-        navigationController?.popViewController(animated: true)
+      let startTime = CACurrentMediaTime()
+      guard let floatPointer = strongSelf.inferencer.module.classify(frames: &pixelBuffer) else {
+        return
+      }
+      let inferenceTime = CACurrentMediaTime() - startTime
+
+      let classes = strongSelf.inferencer.classes
+      DispatchQueue.main.async {
+        let bufferPointer = UnsafeBufferPointer(start: floatPointer, count: classes.count)
+        let zippedResults = zip(bufferPointer, classes.indices)
+        let sortedResults = zippedResults.sorted { $0.0 > $1.0 }.prefix(classes.count).prefix(5)
+        let results = sortedResults.map { classes[$0.1] }
+        strongSelf.lblResult.text = results.joined(separator: ", ") + " - \(Int(1000*inferenceTime))ms"
+      }
     }
+  }
+
+  override func viewWillAppear(_ animated: Bool) {
+    super.viewWillAppear(animated)
+    navigationController?.setNavigationBarHidden(true, animated: false)
+    cameraController.startSession()
+  }
+
+  override func viewWillDisappear(_ animated: Bool) {
+    super.viewWillDisappear(animated)
+    cameraController.stopSession()
+  }
+
+  @IBAction func onBackClicked(_: Any) {
+    navigationController?.popViewController(animated: true)
+  }
 }
diff --git a/TorchVideo/TorchVideo/VideoClassifier.swift b/TorchVideo/TorchVideo/VideoClassifier.swift
index c4add88..ac39360 100644
--- a/TorchVideo/TorchVideo/VideoClassifier.swift
+++ b/TorchVideo/TorchVideo/VideoClassifier.swift
@@ -7,21 +7,21 @@
 import UIKit
 
 class VideoClassifier {
-    lazy var module: InferenceModule = {
-        if let filePath = Bundle.main.path(forResource: "video_classification", ofType: "ptl"),
-            let module = InferenceModule(fileAtPath: filePath) {
-            return module
-        } else {
-            fatalError("Failed to load model!")
-        }
-    }()
-    
-    lazy var classes: [String] = {
-        if let filePath = Bundle.main.path(forResource: "classes", ofType: "txt"),
-            let classes = try? String(contentsOfFile: filePath) {
-            return classes.components(separatedBy: .newlines)
-        } else {
-            fatalError("classes file was not found.")
-        }
-    }()
+  lazy var module: InferenceModule = {
+    guard let filePath = Bundle.main.path(forResource: "video_classification", ofType: "ptl"),
+          let module = InferenceModule(fileAtPath: filePath)
+    else {
+      fatalError("Failed to load model!")
+    }
+    return module
+  }()
+
+  lazy var classes: [Substring] = {
+    guard let filePath = Bundle.main.path(forResource: "classes", ofType: "txt"),
+          let classes = try? String(contentsOfFile: filePath)
+    else {
+      fatalError("classes file was not found.")
+    }
+    return classes.split(whereSeparator: \.isNewline)
+  }()
 }
diff --git a/TorchVideo/TorchVideo/ViewController.swift b/TorchVideo/TorchVideo/ViewController.swift
index ea0a9d0..60dc757 100644
--- a/TorchVideo/TorchVideo/ViewController.swift
+++ b/TorchVideo/TorchVideo/ViewController.swift
@@ -10,151 +10,152 @@ import AVKit
 import AVFoundation
 
 class ViewController: UIViewController, UIImagePickerControllerDelegate, UINavigationControllerDelegate {
-    @IBOutlet weak var btnTest: UIButton!
-    
-    private let testVideos = ["video1", "video2", "video3"]
-    private var videoIndex = 0
-    
-    private var inferencer = VideoClassifier()
-    
-    private var lblResult: UILabel!
-    private var player : AVPlayer?
-    private var playerController :AVPlayerViewController?
-    
-    private var timeObserverToken: Any?
-
-    override func viewWillDisappear(_ animated: Bool) {
-        super.viewWillDisappear(animated)
-
-        player?.pause()
-        if let pc = playerController {
-            lblResult.removeFromSuperview()
-            pc.view.removeFromSuperview()
-            pc.removeFromParent()
-            if let tmObserver = timeObserverToken {
-                player!.removeTimeObserver(tmObserver as Any)
-                timeObserverToken = nil
-            }
-        }
+  @IBOutlet weak var btnTest: UIButton!
+
+  private let testVideos = ["video1", "video2", "video3"]
+  private var videoIndex = 0
+
+  private var inferencer = VideoClassifier()
+
+  private var lblResult: UILabel!
+  private var player : AVPlayer?
+  private var playerController :AVPlayerViewController?
+
+  private var timeObserverToken: Any?
+
+  override func viewWillDisappear(_ animated: Bool) {
+    super.viewWillDisappear(animated)
+
+    player?.pause()
+    if let pc = playerController {
+      lblResult.removeFromSuperview()
+      pc.view.removeFromSuperview()
+      pc.removeFromParent()
+      if let tmObserver = timeObserverToken {
+        player!.removeTimeObserver(tmObserver as Any)
+        timeObserverToken = nil
+      }
     }
-    
-    
-    @IBAction func testTapped(_ sender: Any) {
-        btnTest.setTitle(String(format: "Test %d/%d", videoIndex + 1, testVideos.count), for:.normal)
+  }
 
-        guard let path = Bundle.main.path(forResource: testVideos[videoIndex], ofType:"mp4") else {
-            return
-        }
-        playVideo(url: URL(fileURLWithPath: path), path: path)        
-        videoIndex = (videoIndex + 1) % testVideos.count
+
+  @IBAction func testTapped(_ sender: Any) {
+    btnTest.setTitle(String(format: "Test %d/%d", videoIndex + 1, testVideos.count), for:.normal)
+
+    guard let path = Bundle.main.path(forResource: testVideos[videoIndex], ofType:"mp4") else {
+      return
     }
-    
-    @IBAction func selectTapped(_ sender: Any) {
-        player?.pause()
-        
-        let imagePickerController = UIImagePickerController()
-        imagePickerController.delegate = self;
-        imagePickerController.sourceType = .photoLibrary
-        imagePickerController.mediaTypes = ["public.movie"]
-        self.present(imagePickerController, animated: true, completion: nil)
+    playVideo(url: URL(fileURLWithPath: path), path: path)
+    videoIndex = (videoIndex + 1) % testVideos.count
+  }
+
+  @IBAction func selectTapped(_ sender: Any) {
+    player?.pause()
+
+    let imagePickerController = UIImagePickerController()
+    imagePickerController.delegate = self;
+    imagePickerController.sourceType = .photoLibrary
+    imagePickerController.mediaTypes = ["public.movie"]
+    self.present(imagePickerController, animated: true, completion: nil)
+  }
+
+  @IBAction func liveTapped(_ sender: Any) {
+    player?.pause()
+
+    if UIImagePickerController.isSourceTypeAvailable(.camera) {
+      let imagePickerController = UIImagePickerController()
+      imagePickerController.delegate = self;
+      imagePickerController.sourceType = .camera
+      self.present(imagePickerController, animated: true, completion: nil)
     }
-    
-    @IBAction func liveTapped(_ sender: Any) {
-        player?.pause()
-        
-        if UIImagePickerController.isSourceTypeAvailable(.camera) {
-            let imagePickerController = UIImagePickerController()
-            imagePickerController.delegate = self;
-            imagePickerController.sourceType = .camera
-            self.present(imagePickerController, animated: true, completion: nil)
-        }
+  }
+
+
+  private func playVideo(url: URL, path: String) {
+    if let pc = playerController {
+      lblResult.removeFromSuperview()
+      pc.view.removeFromSuperview()
+      pc.removeFromParent()
+      if let tmObserver = timeObserverToken {
+        player!.removeTimeObserver(tmObserver as Any)
+        timeObserverToken = nil
+      }
     }
-    
-    
-    private func playVideo(url: URL, path: String) {
-        if let pc = playerController {
-            lblResult.removeFromSuperview()
-            pc.view.removeFromSuperview()
-            pc.removeFromParent()
-            if let tmObserver = timeObserverToken {
-                player!.removeTimeObserver(tmObserver as Any)
-                timeObserverToken = nil
-            }
+
+    player = AVPlayer(url: url)
+    playerController = AVPlayerViewController()
+    playerController?.player = player
+
+    playerController?.view.frame = CGRect(x: 0, y: 200, width: 400, height: 240)
+    lblResult = UILabel()
+    lblResult.frame = CGRect(x: 0, y: 200 - 46, width: 400, height: 46)
+    lblResult.backgroundColor = UIColor.blue
+    lblResult.textColor = UIColor.white
+    lblResult.numberOfLines = 2
+    lblResult.lineBreakMode = .byWordWrapping
+    lblResult.textAlignment = .center
+    self.view.addSubview((playerController?.view)!)
+    self.view.addSubview(lblResult)
+    self.addChild(playerController!)
+
+    let timeScale = CMTimeScale(NSEC_PER_SEC)
+    let time = CMTime(seconds: 1.0, preferredTimescale: timeScale)
+
+    timeObserverToken = player!.addPeriodicTimeObserver(forInterval: time, queue: .main) { [weak self] time in
+      guard let self = self else { return }
+
+      DispatchQueue.global().async {
+        let startTime = CACurrentMediaTime()
+        var pixelBuffer = [Float32]()
+        for n in 0...3 {
+          if let image = self.frameFromVideo(path: path, at: time.seconds + Double(n) * 0.3) {
+            let resizedImage = image.resized(to: CGSize(width: CGFloat(Constants.inputWidth), height: CGFloat(Constants.inputHeight)))
+            guard let frameBuffer = resizedImage.normalized() else { return }
+            pixelBuffer += frameBuffer
+          }
+        }
+
+        guard let floatPointer = self.inferencer.module.classify(frames: &pixelBuffer) else { return }
+        let inferenceTime = CACurrentMediaTime() - startTime
+        let classes = self.inferencer.classes
+
+        DispatchQueue.main.async {
+          let bufferPointer = UnsafeBufferPointer(start: floatPointer, count: classes.count)
+          let zippedResults = zip(bufferPointer, classes.indices)
+          let sortedResults = zippedResults.sorted { $0.0 > $1.0 }.prefix(classes.count).prefix(5)
+          let results = sortedResults.map { classes[$0.1] }
+          self.lblResult.text = "\(Int(time.seconds))s: " + results.joined(separator: ", ") + " - \(Int(1000*inferenceTime))ms"
         }
-        
-        player = AVPlayer(url: url)
-        playerController = AVPlayerViewController()
-        playerController?.player = player
-        
-        playerController?.view.frame = CGRect(x: 0, y: 200, width: 400, height: 240)
-        lblResult = UILabel()
-        lblResult.frame = CGRect(x: 0, y: 200 - 46, width: 400, height: 46)
-        lblResult.backgroundColor = UIColor.blue
-        lblResult.textColor = UIColor.white
-        lblResult.numberOfLines = 2
-        lblResult.lineBreakMode = .byWordWrapping
-        lblResult.textAlignment = .center
-        self.view.addSubview((playerController?.view)!)
-        self.view.addSubview(lblResult)
-        self.addChild(playerController!)
-
-        let timeScale = CMTimeScale(NSEC_PER_SEC)
-        let time = CMTime(seconds: 1.0, preferredTimescale: timeScale)
-
-        timeObserverToken = player!.addPeriodicTimeObserver(forInterval: time,
-                                                          queue: .main) {
-            [weak self] time in
-                DispatchQueue.global().async {
-                    let startTime = CACurrentMediaTime()
-                    var pixelBuffer = [Float32]()
-                    for n in 0...3 {
-                        if let image = self!.frameFromVideo(path: path, at: time.seconds + Double(n) * 0.3) {
-                            let resizedImage = image.resized(to: CGSize(width: CGFloat(Constants.inputWidth), height: CGFloat(Constants.inputHeight)))
-                        
-                            guard let frameBuffer = resizedImage.normalized() else { return }
-                            pixelBuffer += frameBuffer
-                        }
-                    }
-
-                    guard let top5Indexes = self!.inferencer.module.classify(frames: &pixelBuffer) else {
-                        return
-                    }
-                    let inferenceTime = CACurrentMediaTime() - startTime
-                                                
-                    DispatchQueue.main.async {
-                        let results = top5Indexes.map { self!.inferencer.classes[$0.intValue] }
-                        self!.lblResult.text = "\(Int(time.seconds))s: " +    results.joined(separator: ", ") + " - \(Int(1000*inferenceTime))ms"
-                    }
-                }
-            }
-
-        player!.play()
+      }
     }
 
-    func frameFromVideo(path: String, at time: Double) -> UIImage? {
-        let asset = AVURLAsset(url: URL(fileURLWithPath: path))
-        let assetIG = AVAssetImageGenerator(asset: asset)
-        assetIG.appliesPreferredTrackTransform = true
-        assetIG.apertureMode = AVAssetImageGenerator.ApertureMode.encodedPixels
+    player!.play()
+  }
 
-        let cmTime = CMTime(value: CMTimeValue(time * 1000000), timescale: CMTimeScale(USEC_PER_SEC))
+  func frameFromVideo(path: String, at time: Double) -> UIImage? {
+    let asset = AVURLAsset(url: URL(fileURLWithPath: path))
+    let assetIG = AVAssetImageGenerator(asset: asset)
+    assetIG.appliesPreferredTrackTransform = true
+    assetIG.apertureMode = AVAssetImageGenerator.ApertureMode.encodedPixels
 
-        let frameRef: CGImage
-        do {
-            frameRef = try assetIG.copyCGImage(at: cmTime, actualTime: nil)
-        } catch let error {
-            print("Error: \(error)")
-            return nil
-        }
+    let cmTime = CMTime(value: CMTimeValue(time * 1000000), timescale: CMTimeScale(USEC_PER_SEC))
 
-        return UIImage(cgImage: frameRef)
+    let frameRef: CGImage
+    do {
+      frameRef = try assetIG.copyCGImage(at: cmTime, actualTime: nil)
+    } catch let error {
+      print("Error: \(error)")
+      return nil
     }
+
+    return UIImage(cgImage: frameRef)
+  }
+
+  func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
+    let videoURL = info[UIImagePickerController.InfoKey.mediaURL] as! URL
+    playVideo(url: videoURL as URL, path: videoURL.path)
+    self.dismiss(animated: true, completion: nil)
     
-    func imagePickerController(_ picker: UIImagePickerController, didFinishPickingMediaWithInfo info: [UIImagePickerController.InfoKey : Any]) {
-        let videoURL = info[UIImagePickerController.InfoKey.mediaURL] as! URL
-        playVideo(url: videoURL as URL, path: videoURL.path)
-        self.dismiss(animated: true, completion: nil)
-    
-    }    
+  }
 }
 
diff --git a/TorchVideo/TorchVideo/video_classification.ptl b/TorchVideo/TorchVideo/video_classification.ptl
new file mode 100644
index 0000000..6f5b510
Binary files /dev/null and b/TorchVideo/TorchVideo/video_classification.ptl differ
diff --git a/TorchVideo/video_classification.ptl b/TorchVideo/video_classification.ptl
new file mode 100644
index 0000000..6f5b510
Binary files /dev/null and b/TorchVideo/video_classification.ptl differ