diff --git a/ImageSegmentation/app/src/main/assets/dog_landscape.jpg b/ImageSegmentation/app/src/main/assets/dog_landscape.jpg
new file mode 100644
index 00000000..615637fb
Binary files /dev/null and b/ImageSegmentation/app/src/main/assets/dog_landscape.jpg differ
diff --git a/ImageSegmentation/app/src/main/java/org/pytorch/imagesegmentation/MainActivity.java b/ImageSegmentation/app/src/main/java/org/pytorch/imagesegmentation/MainActivity.java
index 40222178..1eecc72f 100644
--- a/ImageSegmentation/app/src/main/java/org/pytorch/imagesegmentation/MainActivity.java
+++ b/ImageSegmentation/app/src/main/java/org/pytorch/imagesegmentation/MainActivity.java
@@ -78,8 +78,10 @@ protected void onCreate(Bundle savedInstanceState) {
         final Button buttonRestart = findViewById(R.id.restartButton);
         buttonRestart.setOnClickListener(new View.OnClickListener() {
             public void onClick(View v) {
-                if (mImagename == "deeplab.jpg")
+                if (mImagename.equals("deeplab.jpg"))
                     mImagename = "dog.jpg";
+                else if (mImagename.equals("dog.jpg"))
+                    mImagename = "dog_landscape.jpg";
                 else
                     mImagename = "deeplab.jpg";
                 try {
@@ -130,8 +132,8 @@ public void run() {
         int width = mBitmap.getWidth();
         int height = mBitmap.getHeight();
         int[] intValues = new int[width * height];
-        for (int j = 0; j < width; j++) {
-            for (int k = 0; k < height; k++) {
+        for (int j = 0; j < height; j++) {
+            for (int k = 0; k < width; k++) {
                 int maxi = 0, maxj = 0, maxk = 0;
                 double maxnum = -Double.MAX_VALUE;
                 for (int i = 0; i < CLASSNUM; i++) {
diff --git a/ImageSegmentation/build.gradle b/ImageSegmentation/build.gradle
index 6754c23d..04baf1e7 100644
--- a/ImageSegmentation/build.gradle
+++ b/ImageSegmentation/build.gradle
@@ -5,7 +5,7 @@ buildscript {
         jcenter()
     }
     dependencies {
-        classpath "com.android.tools.build:gradle:4.0.1"
+        classpath 'com.android.tools.build:gradle:4.1.1'
 
         // NOTE: Do not place your application dependencies here; they belong
         // in the individual module build.gradle files
diff --git a/ImageSegmentation/gradle/wrapper/gradle-wrapper.properties b/ImageSegmentation/gradle/wrapper/gradle-wrapper.properties
index 8eba0e70..db97decb 100644
--- a/ImageSegmentation/gradle/wrapper/gradle-wrapper.properties
+++ b/ImageSegmentation/gradle/wrapper/gradle-wrapper.properties
@@ -3,4 +3,4 @@ distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-6.1.1-all.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-6.5-all.zip