diff --git a/build.gradle b/build.gradle
index 8798aba..365318c 100644
--- a/build.gradle
+++ b/build.gradle
@@ -3,7 +3,7 @@ plugins {
   // To create a shadow/fat jar, including dependencies
   id 'com.github.johnrengelman.shadow' version '7.0.0'
   // To manage included native libraries
-  id 'org.bytedeco.gradle-javacpp-platform' version '1.5.6'
+  id 'org.bytedeco.gradle-javacpp-platform' version '1.5.7'
 }
 
 repositories {
@@ -26,7 +26,7 @@ ext.moduleName = 'qupath.extension.stardist'
 
 description = 'QuPath extension to use StarDist'
 
-version = "0.3.0"
+version = "0.3.1-SNAPSHOT"
 
 dependencies {
     def qupathVersion = "0.3.0" // For now
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index 0f80bbf..aa991fc 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,5 +1,5 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-7.0.2-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-7.4.2-bin.zip
 zipStoreBase=GRADLE_USER_HOME
 zipStorePath=wrapper/dists
diff --git a/src/main/java/qupath/ext/stardist/StarDist2D.java b/src/main/java/qupath/ext/stardist/StarDist2D.java
index d443345..9b88a40 100644
--- a/src/main/java/qupath/ext/stardist/StarDist2D.java
+++ b/src/main/java/qupath/ext/stardist/StarDist2D.java
@@ -17,6 +17,7 @@
 package qupath.ext.stardist;
 
 import java.awt.image.BufferedImage;
+import java.io.Closeable;
 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
@@ -47,8 +48,10 @@
 import org.locationtech.jts.geom.Coordinate;
 import org.locationtech.jts.geom.Envelope;
 import org.locationtech.jts.geom.Geometry;
+import org.locationtech.jts.geom.GeometryCollection;
 import org.locationtech.jts.geom.Location;
 import org.locationtech.jts.geom.Polygon;
+import org.locationtech.jts.geom.util.GeometryFixer;
 import org.locationtech.jts.index.strtree.STRtree;
 import org.locationtech.jts.simplify.VWSimplifier;
 import org.slf4j.Logger;
@@ -97,7 +100,7 @@
  * 
  * @author Pete Bankhead (this implementation, but based on the others)
  */
-public class StarDist2D {
+public class StarDist2D implements AutoCloseable {
 
 	private final static Logger logger = LoggerFactory.getLogger(StarDist2D.class);
 	
@@ -537,7 +540,7 @@ public Builder inputAdd(double... values) {
 		 * @return this builder
 		 */
 		public Builder inputScale(double... values) {
-			this.ops.add(ImageOps.Core.subtract(values));
+			this.ops.add(ImageOps.Core.multiply(values));
 			return this;
 		}
 		
@@ -921,11 +924,25 @@ private PathObject convertToObject(PotentialNucleus nucleus, ImagePlane plane, d
 		var geomNucleus = simplify(nucleus.geometry);
 		PathObject pathObject;
 		if (cellExpansion > 0) {
+//			cellExpansion = geomNucleus.getPrecisionModel().makePrecise(cellExpansion);
+//			cellExpansion = Math.round(cellExpansion);
 			var geomCell = CellTools.estimateCellBoundary(geomNucleus, cellExpansion, cellConstrainScale);
-			if (mask != null)
+			if (geomCell instanceof GeometryCollection && geomNucleus instanceof Polygon) {
+				if (!geomCell.isValid()) {
+					// Sometimes buffer creates invalid geometries
+					// (Note that the fix should already be applied by estimateCellBoundary in v0.4.0)
+					geomCell = GeometryFixer.fix(geomCell);
+					logger.debug("Used GeometryFixer to fix an invalid cell boundary geometry");
+				}
+			}
+			if (mask != null) {
 				geomCell = GeometryTools.attemptOperation(geomCell, g -> g.intersection(mask));
+			}
 			geomCell = simplify(geomCell);
 			
+			// Intersection with complex mask could give linestrings - which we want to remove
+			geomCell = GeometryTools.ensurePolygonal(geomCell);
+			
 			if (geomCell.isEmpty()) {
 				logger.warn("Empty cell boundary at {} will be skipped", nucleus.geometry.getCentroid());
 				return null;
@@ -1069,6 +1086,13 @@ private List<PotentialNucleus> detectObjectsForTile(ImageDataOp op, DnnModel<?>
 			requestPadded = RegionRequest.createInstance(server.getPath(), downsample, x1, y1, x2-x1, y2-y1);
 		}
 		
+//		// Hack to visualize the tiles that are computed (for debugging)
+//		imageData.getHierarchy().addPathObject(
+//				PathObjects.createAnnotationObject(
+//						ROIs.createRectangleROI(request),
+//						PathClassFactory.getPathClass("Temporary")
+//						));
+		
 		try (@SuppressWarnings("unchecked")
 		var scope = new PointerScope()) {
 			Mat mat;
@@ -1350,6 +1374,11 @@ private List<PotentialNucleus> filterNuclei(List<PotentialNucleus> potentialNucl
 	                if (envelope.intersects(env) && nucleus.geometry.intersects(nucleus2.geometry)) {
 	                	// Retain the nucleus only if it is not fragmented, or less than half its original area
 	                    var difference = nucleus2.geometry.difference(nucleus.geometry);
+	                    
+	                    // Discard linestrings
+	                    if (difference instanceof GeometryCollection)
+	                    	difference = GeometryTools.ensurePolygonal(difference);
+	                    
 	                    if (difference instanceof Polygon && difference.getArea() > nucleus2.fullArea / 2.0)
 	                        nucleus2.geometry = difference;
 	                    else {
@@ -1406,5 +1435,21 @@ int getClassification() {
 	    }
 		
 	}
+
+
+	/**
+	 * Close and cleanup resources.
+	 * 
+	 * @implNote In practice, this means close any {@link DnnModel} stored if it is an instance of 
+	 * {@link Closeable} or {@link AutoCloseable}.
+	 * This can be important to avoid memory leaks, particularly if using a GPU.
+	 */
+	@Override
+	public void close() throws Exception {
+		if (dnn instanceof Closeable) {
+			((Closeable) dnn).close();
+		} else if (dnn instanceof AutoCloseable)
+			((AutoCloseable) dnn).close();
+	}
 	
 }
diff --git a/src/main/java/qupath/ext/stardist/StarDistExtension.java b/src/main/java/qupath/ext/stardist/StarDistExtension.java
index f18db05..8687425 100644
--- a/src/main/java/qupath/ext/stardist/StarDistExtension.java
+++ b/src/main/java/qupath/ext/stardist/StarDistExtension.java
@@ -51,7 +51,7 @@ public String getDescription() {
 	
 	@Override
 	public Version getQuPathVersion() {
-		return Version.parse("0.3.0-rc2");
+		return Version.parse("0.3.0");
 	}
 
 	@Override