| 
					
				 | 
			
			
				@@ -129,10 +129,12 @@ def allow_oom(cls): 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				         def _wrapper(self, *args, **kwargs): 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				             try: 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				                 func(self, *args, **kwargs) 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				-            except (SystemError, RuntimeError, OSError) as e: 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				+            except (SystemError, RuntimeError, OSError, MemoryError) as e: 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				+                # XXX: This may not cover all OOM cases. 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				                 msg = str(e) 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				                 if "Out of memory error" in msg \ 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				-                    or "(External) CUDNN error(4), CUDNN_STATUS_INTERNAL_ERROR." in msg: 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				+                    or "(External) CUDNN error(4), CUDNN_STATUS_INTERNAL_ERROR." in msg \ 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				+                    or isinstance(e, MemoryError): 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				                     logging.warning("An OOM error has been ignored.") 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				                 else: 
			 | 
		
	
		
			
				 | 
				 | 
			
			
				                     raise 
			 |