diff --git a/judge/highlight_code.py b/judge/highlight_code.py
index ccef8d4..979eebf 100644
--- a/judge/highlight_code.py
+++ b/judge/highlight_code.py
@@ -7,17 +7,10 @@ def _make_pre_code(code):
     return mark_safe("<pre>" + escape(code) + "</pre>")
 
 
-def _wrap_code(inner):
-    yield 0, "<code>"
-    for tup in inner:
-        yield tup
-    yield 0, "</code>"
-
-
 try:
     import pygments
     import pygments.lexers
-    import pygments.formatters.html
+    import pygments.formatters
     import pygments.util
 except ImportError:
 
@@ -25,11 +18,6 @@ except ImportError:
         return _make_pre_code(code)
 
 else:
-
-    class HtmlCodeFormatter(pygments.formatters.HtmlFormatter):
-        def wrap(self, source, outfile):
-            return self._wrap_div(self._wrap_pre(_wrap_code(source)))
-
     def highlight_code(code, language, cssclass="codehilite", linenos=True):
         try:
             lexer = pygments.lexers.get_lexer_by_name(language)
@@ -39,9 +27,9 @@ else:
         if linenos:
             return mark_safe(
                 pygments.highlight(
-                    code, lexer, HtmlCodeFormatter(cssclass=cssclass, linenos="table")
+                    code, lexer, HtmlCodeFormatter(cssclass=cssclass, linenos="table", wrapcode=True)
                 )
             )
         return mark_safe(
-            pygments.highlight(code, lexer, HtmlCodeFormatter(cssclass=cssclass))
+            pygments.highlight(code, lexer, HtmlCodeFormatter(cssclass=cssclass, wrapcode=True))
         )