diff --git a/TensorFI/fiConfig.py b/TensorFI/fiConfig.py
index 4554caa..fe955c6 100644
--- a/TensorFI/fiConfig.py
+++ b/TensorFI/fiConfig.py
@@ -171,10 +171,11 @@ def __str__(self):
 		res.append(" }")
 		return "\n".join(res)
 
-	def __init__(self,fiParams):
+	def __init__(self, fiParams, opProbabilities):
 		"Configure the initial fault injection parameters from the fiParams Dictionary"
 		# First configure the Scalar fault type
 		# Default value of fault is NoFault
+		self.opProbabilities = opProbabilities
 		if fiParams.has_key(Fields.ScalarFaultType.value):
 			faultTypeScalar = fiParams[Fields.ScalarFaultType.value]
 		else:	
diff --git a/TensorFI/injectFault.py b/TensorFI/injectFault.py
index f7ccc15..56bcc3f 100644
--- a/TensorFI/injectFault.py
+++ b/TensorFI/injectFault.py
@@ -23,7 +23,7 @@
 # and is called from TensorFI.py's constructor
 
 # NOTE: This has to be in this module or else fiConf won't be accessible
-def initFIConfig(fiParams):
+def initFIConfig(fiParams, opProbabilities = []):
 	"Initialize the global variable fiConf with the params"
 	global fiConf
 	global count
@@ -36,7 +36,7 @@ def initFIConfig(fiParams):
 	# which op to be injected in the whole run
 	global injectedOp
 
-	fiConf = FIConfig(fiParams)
+	fiConf = FIConfig(fiParams, opProbabilities)
 	logging.debug("Initialized config file : " + str(fiConf))
 	
 	# Setup the random seed for the fault injector if one is specified
@@ -234,7 +234,7 @@ def condPerturb(op, res):
 				totalVistedOp += 1
 				# select one random op to be injected in the whole run
 				if(injectedOp == 0):
-					injectedOp = np.random.randint(low=1, high=totalInstance+1) 
+					injectedOp = np.random.choice(totalInstance, 1, fiConf.opProbabilities) + 1
 				# inject fault at the output of the operation
 				if(totalVistedOp == injectedOp):
 					res = perturb(res)
diff --git a/TensorFI/tensorFI.py b/TensorFI/tensorFI.py
index b862760..7ee1bc4 100644
--- a/TensorFI/tensorFI.py
+++ b/TensorFI/tensorFI.py
@@ -133,6 +133,17 @@ def printInjectMap(self):
 			print "\t", tensor.name, " : ", fiTensor.name
 		print "Done fiMap"
 
+	def setProbability(self):
+		"Set the probabilities for FI in operators based on the output state space"
+		opProbability = []
+		totalStateSpace = 0.
+		for (tensor, fiTensor) in self.fiMap.iteritems():
+			opShape = tensor.shape.num_elements()
+			opProbability.append(opShape)
+			totalStateSpace += opShape
+		opProbability[:] = [x / totalStateSpace for x in opProbability]
+		return opProbability
+
 	# These are all functions that can be called externally for interacting with the injector
 
 	def turnOffInjections(self):
@@ -193,7 +204,13 @@ def __init__(self, s,	# This is the session from tensorFlow
 		# Configure the fault injection parameters for the injection functions	
 		# fiConf is a global variable as it needs to be accessible to the FI functions
 		logging.info("Initializing the fault injection parameters")
-		initFIConfig(fiParams)
+
+		# Random sampling across state space distribution for one fault per run
+		if fiParams["InjectMode"] == "oneFaultPerRun":
+			opProbabilities = self.setProbability()
+			initFIConfig(fiParams, opProbabilities)
+		else:
+			initFIConfig(fiParams)
 
 		# Initialize the default fault log - this may be overridden by the launch method later
 		# This is in case the run method is called directly without going through the launch